[llama4 -> llama4_jax] Refactor to be a proper installable Python package #9

SamuelMarks · 2025-04-17T17:03:03Z

[llama4 -> llama4_jax] Refactor to be a proper installable Python package ; [llama4_jax/pyproject.toml] Add missing dependency ; [llama4_jax/README.md] Document new usage

Also I know you have path = Path(path).expanduser().absolute() but that doesn't provide nice --help text and should be expanded earlier anyway. I can remove your Path(path).expanduser() if you give the go-ahead.

PS: If you merge this PR I can send you a new PR for deepseek_r1_jax

…kage ; [llama4_jax/pyproject.toml] Add missing dependency ; [llama4_jax/README.md] Document new usage

rdyro

Thanks for this PR! I left some comments

rdyro · 2025-04-24T00:54:36Z

llama4_jax/scripts/convert_weights.py

    additional_files = ["config.json", "tokenizer.json", "tokenizer_config.json"]
    for additional_file in additional_files:
-        full_paths = list(model_path.glob(f"**/{additional_file}"))
+        full_paths = list(model_path.glob(os.path.join("**", additional_file)))


can we stick with pathlib.Path throughout here, it largely replaces os.path.*

After 11 years maybe I can rely on it being available!

os.path.join is actually used by it internally IIRC; but sure can move to pathlib for my jax-ml/jax-llm-examples contributions.

rdyro · 2025-04-24T00:55:52Z

llama4_jax/scripts/download_model.py

        "--dest-root-path",
        required=True,
-        default="~/",
+        default=os.path.join(os.path.expanduser("~"), ""),


Path("~/").expanduser()

rdyro · 2025-04-24T00:56:19Z

llama4_jax/scripts/quantize_model.py

    parser = ArgumentParser()
    parser.add_argument(
-        "--path", default="~/DeepSeek-R1-Distill-Llama-70B", required=True, help="Existing JAX model checkpoint path"
+        "--path", default=os.path.join(os.path.expanduser("~"), "DeepSeek-R1-Distill-Llama-70B"),


Path("~/DeepSeek-R1-Distill-Llama-70B").expanduser()

rdyro · 2025-04-24T00:56:36Z

llama4_jax/pyproject.toml

 name = "llama4_jax"
 version = "0.1.0"
-description = ""
+description = "Pure JAX implementation of Llama 4 inference, including a checkpoint converter for the weights."


rdyro · 2025-04-24T01:19:18Z

llama4_jax/README.md


-The entire model is defined in [model.py](llama4_jax/model.py) and invoked
-via [main.py](main.py). Among other things, the model code demonstrates:
+The entire model is defined in [__main__.py](llama4_jax/__main__.py) and invoked


Can we keep main.py separate, I think an explicit script might be more clear that it's just a starting point for a larger program rather than default behavior for the llama4 module

@rdyro It's your call. But it is very nonstandard Python, the __main__.py properly integrates… seeing itself as within Python module; being relocatable; and installable.

One other idea is to create a hierarchy:

. ├── jax_examples_cli ├── deepseek_r1_jax └── llama4_jax

All packages are installable. There are no main.pys. Outside of jax_examples_cli there are no __main__.pys. In the Command Line Interface it finds installed packages and/or packages within a specific location (e.g., os.getcwd() or JAX_LLM_EXAMPLES env var); indicating where to find packages compatible with the jax_examples_cli/__main__.py.

Already your main.pys are very similar, it wouldn't be hard to hoist them up. Usage would be: jax_examples_cli --model <deepseek-r1-jax | llama4_jax> --ckpt_path …

# Conflicts: # llama4_jax/llama4_jax/model.py # llama4_jax/scripts/convert_weights.py # llama4_jax/scripts/download_model.py # llama4_jax/scripts/quantize_model.py

…face-hub` to deps ; sort deps

rdyro · 2025-06-24T20:50:58Z

I like the changes! One request, let's please avoid os.path entirely in favor of pathlib.Path, I prefer not to mix them.

Lots of changes, so I'll need to provision some compute to test those, give me a couple of days, thanks!

rdyro

Hey, I really appreciate the desire to make this repo better, but I cannot merge this PR as is. It's introducing too many unrelated changes at once:

import sorting in python files
dependency sorting in pyproject files
extra spaces for import groups separation
renaming llama4 to llama4_jax
renaming of main.py to __main__.py
support for non-unix path separators

Can you please separate those into separate PRs or focus on only one of those?

rdyro · 2025-08-11T19:38:56Z

llama4_jax/pyproject.toml

 license = { text = "Apache-2.0" }

 dependencies = [
+    "datasets",


please undo import sorting

rdyro · 2025-08-11T19:39:16Z

llama4_jax/llama4_jax/ragged_attention.py

 from jax.experimental.pallas import tpu as pltpu
 from jax.experimental.shard_map import shard_map
 from jax.sharding import Mesh, PartitionSpec as P, NamedSharding
+


revert extra spaces in imports

rdyro · 2025-08-11T19:39:19Z

llama4_jax/llama4_jax/model.py

 from jax.experimental.shard_map import shard_map
 from jax.sharding import PartitionSpec as P, use_mesh
 from jax.experimental.shard import auto_axes, reshard
+


revert extra spaces in imports

rdyro · 2025-08-11T19:39:27Z

llama4_jax/llama4_jax/chkpt_utils.py

 import jax
 from jax import numpy as jnp
 from jax.sharding import PartitionSpec as P
+


revert extra spaces in imports

rdyro · 2025-08-11T19:40:34Z

deepseek_r1_jax/deepseek_r1_jax/chkpt_utils.py

 import re
 from concurrent.futures import ThreadPoolExecutor
 import math
+


revert extra spaces in imports

rdyro · 2025-08-11T19:40:53Z

deepseek_r1_jax/deepseek_r1_jax/model.py

 from pathlib import Path

+from etils import epath
+


revert extra spaces in imports

rdyro · 2025-08-11T19:41:35Z

qwen3/scripts/convert_weights.py


    model_path, ckpt_path = Path(model_path).expanduser(), Path(ckpt_path).expanduser()
-    files = list(model_path.glob("**/*safetensors"))
+    files = list(model_path.glob(os.path.join("**", "*safetensors")))


don't mix pathlib and os.path

rdyro · 2025-08-11T19:41:38Z

qwen3/scripts/convert_weights.py

+    files = list(model_path.glob(os.path.join("**", "*safetensors")))
    assert len(files) > 1
-    config_files = list(model_path.glob("**/config.json"))
+    config_files = list(model_path.glob(os.path.join("**", "config.json")))


don't mix pathlib and os.path

rdyro · 2025-08-11T19:41:55Z

qwen3/main.py


 import dataclasses
+import os
+


revert extra spaces in imports

rdyro · 2025-08-11T19:42:25Z

llama4_jax/scripts/convert_weights.py


-def main(model_path: str | Path, ckpt_path: str | Path):


importing at top level makes for a very slow help message, import under main

[llama4 -> llama4_jax] Refactor to be a proper installable Python pac…

6fd621f

…kage ; [llama4_jax/pyproject.toml] Add missing dependency ; [llama4_jax/README.md] Document new usage

rdyro reviewed Apr 24, 2025

View reviewed changes

SamuelMarks added 2 commits June 23, 2025 12:02

Merge branch 'main' into PKG-FIX__llam4_jax

2286f87

# Conflicts: # llama4_jax/llama4_jax/model.py # llama4_jax/scripts/convert_weights.py # llama4_jax/scripts/download_model.py # llama4_jax/scripts/quantize_model.py

[*.py] Use os.path.join / Path ; [**/pyproject.toml] Add `hugging…

d5c936a

…face-hub` to deps ; sort deps

[llama4_jax/scripts/convert_weights.py] Remove debug code

3812791

rdyro requested changes Aug 11, 2025

View reviewed changes

[llama4 -> llama4_jax] Refactor to be a proper installable Python package #9

Are you sure you want to change the base?

[llama4 -> llama4_jax] Refactor to be a proper installable Python package #9

Uh oh!

Conversation

SamuelMarks commented Apr 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rdyro left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rdyro commented Jun 24, 2025

Uh oh!

rdyro left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

SamuelMarks commented Apr 17, 2025 •

edited

Loading