model-bringup-m3

# MaxText M3 Model Bring-Up (`src/maxtext/m3/`) Onboards a HuggingFace model family as a self-contained pure-NNX vertical slice in `src/maxtext/m3/models/<family>/`, verified end-to-end on CPU and TPU via `pytest`. For the motivation and design rationale behind m3 (model garden layout, sharding as annotation, post-construction transforms, MoE runtime split), read `src/maxtext/m3/DESIGN.md`. Where it disagrees with the code or `src/maxtext/m3/README.md`, the code and README are the source of truth. ## 1. Vertical-Slice Isolation & M3 Architecture Rules **Allowed touch list:** - `src/maxtext/m3/models/<family>/`: - `__init__.py` — public exports (`<Family>Config`, `<Family>Sharding`, `<Family>Model`, `create_<family>_model`, helper and mapping exports) - `modeling_<family>.py` — config, sharding rules, norms, attention, MLP/MoE, stacked decoder, model, loss - `checkpoint_mappings_<family>.py` — declarative MaxText <-> HuggingFace parameter mapping (`<FAMILY>_MAXTEXT_TO_HF_PARAM_MAPPING`), tensor transform hooks (`<FAMILY>_MAXTEXT_TO_HF_PARAM_HOOK_FN`), and HF tensor shapes (`<FAMILY>_HF_WEIGHTS_TO_SHAPE`), plus lightweight table-driven in-memory weight helpers for unit tests (reusing the centralized conversion planner rather than a bespoke per-model CLI converter) - `utils_<family>.py` — `calculate_<family>_tflops_per_device(config, log=True)` and optional model-related helpers such as `configure_<family>_engine(engine)` (`MaxTextTrainingEngine` `aux_loss` adapter) - `src/maxtext/m3/models/__init__.py` — lazy wrappers in `MODEL_REGISTRY`, `MODEL_FLOPS_REGISTRY`, `MODEL_ENGINE_CONFIG_REGISTRY` - **Training Entrypoint Reuse (`src/maxtext/m3/train/loop.py`):** Text-only LLMs and MoEs reuse M3's dedicated training entrypoint `src/maxtext/m3/train/loop.py` (`train_loop` / `train_step`) directly — do **not** add a specialized `<family>_train.py` trainer. Only introduce a thin shim over `loop.py` when non-text/multimodal input or output modalities require model-specific data/step handling. - **Centralized Checkpoint Conversion & Config Registration:** - `src/maxtext/configs/models/<model-name>.yml` — set `use_m3_model: true` and `param_scan_axis: 0` (omitting `"model_type"` so `m3.configs.config` selects `BaseConfig`, and keeping fields compatible with `MaxTextConfig` so both `pyconfig.initialize` and `m3.configs.config.initialize` load cleanly) - `src/maxtext/configs/types.py` & `src/maxtext/utils/globals.py` — register `<model-name>` in `ModelName` (enforcing `use_m3_model=True`) and `HF_IDS` - `src/maxtext/checkpoint_conversion/utils/{param_mapping.py,hf_shape.py,hf_model_configs.py}` — wire `checkpoint_mappings_<family>` into `PARAM_MAPPING`, `HOOK_FNS`, `HF_SHAPE`, and `HF_MODEL_CONFIGS` so `to_maxtext` and `to_huggingface` convert checkpoints end-to-end - `tests/unit/m3/<family>*_test.py` — conformance, `PerturbSpec`, centralized checkpoint mapping/roundtrip, PyTorch parity, and TPU verification tests **Never modify** legacy `src/maxtext/layers/`, `src/maxtext/models/`, or `model_creation_utils.py`. (Note: `.agents/` is listed in `.gitignore`; stage new skill files with `git add -f` in OSS.) **Seven Load-Bearing NNX & `maxtext_engine` Invariants (`src/maxtext/m3/README.md`):** 1. **Pure Flax NNX built-ins:** Compose `nnx.Linear`, `nnx.Embed`, `nnx.RMSNorm`, `nnx.Param`, and `m3.core.rope.apply_rope` directly. No Linen bridges, no shared `Decoder`, no model-name branching. 2. **`nnx.Param` is the sole trainable variable type:** Other `nnx.Variable` subclasses silently receive zero gradient in `maxtext_engine`. 3. **Never `sow` or mutate non-`Param` state in `__call__`:** Mutating non-`Param` treedef permanently drops the engine onto its slow path. Pass inference `KVCache` explicitly as an external PyTree argument; default `return_aux: bool = False` in `__call__` (returning `logits` for `m3/train/loop.py` and `(logits, aux_dict)` when `return_aux=True`). 4. **Path-transparent `scan` and `remat`:** Create decoder layers as a single stacked module via `nnx.vmap` (`(num_layers, ...)` leading axis on every `nnx.Param`). Unrolled (`scan_layers=False`) and scanned (`scan_layers=True`) execution MUST share the exact same `nnx.GraphDef` and stacked parameter paths. **Never store Python `int` `self.layer_idx = i` on scanned submodules** (breaks `nnx.split` graphdef equality across layers). 5. **Fold auxiliary losses into primary loss:** Return `aux["moe_lb_loss"]` when `return_aux=True`, and wire `utils_<family>.configure_<family>_engine` into `MODEL_ENGINE_CONFIG_REGISTRY` via `engine.with_loss_fn(..., has_aux=True)` (keep `targets_segmentation` strictly on CE loss — never feed it into self-attention `attention_mask` or `moe_lb_loss`). 6. **Single shared `Mesh` + `nnx.eval_shape`-safe eager sharding:** Accept/validate `quant_mode_str: str | None = None` in `create_<family>_model` without forwarding it to `<Family>Config`. When not tracing (`not isinstance(jnp.zeros(()), jax.core.Tracer)`), construct inside `with jax.sharding.set_mesh(mesh), nnx.use_eager_sharding(True):`; whenever tracing (`nnx.eval_shape` or `jax.jit` inside `create_nnx_sharded_model` / `from_pretrained`), construct inside `with nnx.use_eager_sharding(False):` so abstract and jitted `nnx.Param` states share identical `eager_sharding` PyTree metadata. Set `self.mesh = mesh` (`model.mesh is driver_mesh`) and never store private `Mesh` objects on child modules. 7. **Zero post-init parameter constraints:** Shard weights once at init via `kernel_metadata={"out_sharding": sharding(...)}`. Inside `__call__`, call `sharding.constrain(x, name, axes)` on **activations only** — never call `with_sharding_constraint`, `reshard`, or `device_put` on parameter-derived weight arrays. **Stacked Layer Construction & `NamedSharding` Placement Pattern:** ```python forked_rngs = rngs.split(config.num_hidden_layers) stacked = nnx.vmap(lambda r: FamilyDecoderLayer(config, sharding=sharding, rngs=r), in_axes=0, out_axes=0)(forked_rngs) graphdef, state = nnx.split(stacked) def _place_stacked(path, var): if not isinstance(var, nnx.Variable): return var orig = var.get_metadata("out_sharding", None) orig_spec = orig.spec if isinstance(orig, jax.sharding.NamedSharding) else orig new_spec = jax.sharding.PartitionSpec(None, *orig_spec) raw_val = var.get_raw_value() if mesh is not None and not mesh.empty and not isinstance(raw_val, jax.core.Tracer): return var.replace(jax.device_put(raw_val, jax.sharding.NamedSharding(mesh, new_spec)), out_sharding=new_spec) return var.replace(raw_val, out_sharding=new_spec) layers = nnx.merge(graphdef, jax.tree_util.tree_map_with_path(_place_stacked, state, is_leaf=lambda x: isinstance(x, nnx.Variable))) ``` At forward time, split `graphdef, params, non_params = nnx.split(self.layers, nnx.Param, ...)`, optionally wrap the per-layer body in `jax.checkpoint(..., policy=...)`, and run `jax.lax.scan` when `scan_layers=True` or a Python loop slicing axis `0` when `scan_layers=False`. ## 2. 4-Axis Logical Sharding, MoE Dispatch & Lazy Registries Subclass `maxtext.m3.core.sharding.Sharding` over `MESH_AXIS_NAMES = ("data", "fsdp", "tensor", "expert")`: ```python class FamilySharding(Sharding): def __init__(self, vocab_size: int = 0, num_experts: int = 1, tp_size: int = 1, fsdp_size: int = 1, ep_size: int = 1): self.vocab_tp = (vocab_size > 0 and vocab_size % tp_size == 0) self.compound_ep = (num_experts % (fsdp_size * ep_size) == 0) def map_axis(self, axis: str, tensor_name: str, tensor_type: TensorType) -> MeshAxes: match axis, tensor_type: case "batch", TensorType.ACTIVATION: return ("data", "fsdp", "expert") case "embed", TensorType.WEIGHT: return ("fsdp", "expert") case "experts", TensorType.WEIGHT: return ("fsdp", "expert") if self.compound_ep else "expert" case "embed_exp", TensorType.WEIGHT: return None if self.compound_ep else "fsdp" case ("heads" | "kv_heads" | "mlp"), _: return "tensor" case "vocab", _: return "tensor" if self.vocab_tp else None case ("length" | "embed" | "head_dim" | "norm"), _: return None case _: raise KeyError(f"Unknown logical axis {axis!r} on {tensor_name!r}") ``` - **Non-divisible `vocab_size` guard:** Odd vocabularies (e.g., `49155`) leave `"vocab"` replicated (`None`) when `vocab_size % tp_size != 0` while still sharding `embed_tokens` `(V, H)` along `"embed"` (`("fsdp", "expert")`). - **2-D+ Weight Coverage:** Every 2-D+ weight (`>= 256` unstacked elements) must shard across at least one active axis on `(data=1, fsdp=2, tensor=2, expert=1)` (store router kernel as `(hidden_size, num_experts)` on `("embed", None)`). - **MoE Layout & Zero-All-Gather Dispatch (`FSDP >= 1, TP >= 1, EP >= 1`):** - Store `gate_up_proj` as 4-D `(E, H, 2, I)` (`("experts", "embed_exp", None, "mlp")`) so gate `[:, :, 0, :]` and up `[:, :, 1, :]` slice the replicated size-2 axis; store `down_proj` as 3-D `(E, I, H)` (`("experts", "mlp", "embed_exp")`). - When `compound_ep` holds and `fsdp * expert > 1` (or `tensor > 1`), dispatch inside `jax.shard_map(..., check_vma=False)` with exact parameter `PartitionSpec`s: keep axis tuple order `("fsdp", "expert")` identical across `P(("fsdp", "expert"), ...)` weights, `lax.axis_index(("fsdp", "expert"))` (`f_idx * ep_size + e_idx`), `lax.all_gather(..., ("fsdp", "expert"), axis=0, tiled=True)`, and `lax.psum_scatter(..., ("fsdp", "expert"), scatter_dimension=0, tiled=True)` (plus `lax.psum` over `"tensor"`). When token count `T` is not divisible by `fsdp * expert` (e.g., `B=1` KV-cache decode), fall back activation `in_specs`/`out_specs` to `P(None, None)` + `lax.psum(..., ("tensor", "fsdp", "expert"))` while keeping expert weights sharded. - **TPU `ragged_dot` VJP Tail Masking:** When `sum(group_sizes) < M` (partial local expert slots under EP/FSDP), explicitly mask non-local tail rows `[sum(group_sizes):M]` to `0` on `sorted_x`, SwiGLU `act`, and `sorted_out` via `jnp.where(sorted_valid, ..., 0)` — TPU XLA's `ragged-dot` lhs VJP does not zero rows past `sum(group_sizes)`. - **Lazy Registries & Consolidated Model Helpers (`src/maxtext/m3/models/__init__.py`, `utils_<family>.py`):** - In `MODEL_REGISTRY`, validate `if rngs is None: raise ValueError(...)` and import `<family>` modules lazily inside wrapper functions.