**Link:** https://github.com/pytorch/pytorch/pull/<N> **Summary:** <first ~20 lines of PR body> ### Changed files <git show --stat output> ### Full bisect log <git bisect log output> ``` Omit the PR section if no PR number was found. Do NOT diagnose root cause. ## Rules - Always use absolute paths. - Judge test results by exit code only — never analyze output. - On unexpected errors, show the error and ask the user before proceeding. - Keep the user informed after every bisect step.  --- name: add_moe_model description: Add a new HF MoE model to the transformers_modeling_backend experiment. Use when the user provides a HuggingFace model ID and wants to integrate it with Titan's MoE, or invokes /add_moe_model. --- # Add HF MoE Model to Transformers Backend Add support for a new HuggingFace MoE model in the `torchtitan/experiments/transformers_modeling_backend/` experiment. ## Scope of changes **ONLY modify code under `torchtitan/experiments/transformers_modeling_backend/`.** If supporting the new model appears to require a change outside this folder (e.g. in `torchtitan/models/`, `torchtitan/distributed/`, `torchtitan/components/`, or any other core path), **STOP and escalate to the user** — describe the change you believe is needed and why, and wait for their decision instead of editing core yourself. The one exception is the **temporary, must-be-reverted diagnostic patches** in Phase 3c (numerical-gap analysis): those edit core only to measure where titan diverges from HF, are reverted before the change is done, and are never committed. Any *permanent* change outside the experiment folder still requires escalation. **DO modify experiment code** to support the new model. The experiment folder (`torchtitan/experiments/transformers_modeling_backend/`) is designed to be extended for new models. Expected changes include: - `state_dict_adapter.py` — add regex patterns for new key naming conventions, handle transposed weight layouts, fix round-trip mappings - `moe_replacement.py` — add detection logic for new scoring functions, router types, expert layouts, or shared expert patterns in `_probe_hf_moe_block` and its helper functions - `hf_sharding.py` — give **every parameter/buffer-bearing module** the model introduces a sharding config. This is not limited to attention projections: `set_hf_sharding_configs` covers the root-level `tok_embeddings`, `norm`, `lm_head`, and `rotary_emb` (and their buffers), and `_set_layer_sharding_configs` covers per-layer norms, the attention block, and the dense MLP. A module with params/buffers but no config will mix a plain tensor with a DTensor and crash under TP (the `_assert_all_states_sharded` backstop now makes that a loud setup-time error for any param/buffer-bearing dense-path module — root modules and every decoder layer, excluding the Titan-swapped MoE). For a genuinely new module, shard it (`colwise_sharding_config`/`rowwise_sharding_config`) if it has a clear TP layout, otherwise default to `_replicate_sharding_config(module)` (replicate weights, no TP). If a new module type doesn't fit any existing helper, that's a signal to stop and think (or escalate) rather than leave it unconfigured. - `model.py` — handle alternative submodule names: HF models name the same component differently (feed-forward `mlp` vs `feed_forward`, final norm `norm` vs `final_layernorm`). Add a new model's variant to the relevant lookup (`_get_moe_attr_name`, the `HFTransformerModel` accessor properties). - `scripts/numerical_equivalence.py` — add the model's synthetic test config When blockers are found in Phase 1, fix them in the experiment code before proceeding to Phase 2. Do NOT stop and ask — implement the fixes, then continue the workflow. ## Input The user provides a HuggingFace model ID (e.g. `Qwen/Qwen3-30B-A3B`). Optionally, a local path to downloaded weights for pretrained testing. Before starting, ask the user (if not already specified): 1. **Which branch to work on** — should changes go on the current branch, or a new branch? If running in a worktree, clarify whether the commit should target the main feature branch or stay on the worktree branch. 2. **Whether to commit** — after all tests pass, should the changes be committed automatically, or left as uncommitted changes for the user to review first? ## Phase 1: Compatibility Probe Load the model config and a single MoE layer on meta device to verify titan compatibility. No file changes in this phase — but record all issues found. After the probe, fix any blockers in the experiment code before moving to Phase 2. ### 1a. Load and inspect config ```python from transformers import AutoConfig, AutoModelForCausalLM config = AutoConfig.from_pretrained(model_id, trust_remote_code=True) ``` Record: `model_type`, `hidden_size`, `intermediate_size`, `moe_intermediate_size`, `num_attention_heads`, `num_key_value_heads`, `num_hidden_layers`. ### 1b. Instantiate on meta device ```python with torch.device("meta"): model = AutoModelForCausalLM.from_config(config, trust_remote_code=True) ``` ### 1c. Find MoE layers Scan `model.model.layers` for layers where `hasattr(mlp, "gate") or hasattr(mlp, "router")` AND `hasattr(mlp, "experts")`. Record: - Which layer indices are MoE vs dense - The mechanism: `first_k_dense_replace`, `decoder_sparse_step`, or `mlp_layer_types` - Minimum `num_hidden_layers` needed to have at least one MoE layer ### 1d. Probe MoE block Call `_probe_hf_moe_block(moe_block, config)` from `torchtitan/experiments/transformers_modeling_backend/moe_replacement.py`. **Check each field succeeds without error:** | Field | How it's resolved | Blocker if missing | |-------|-------------------|--------------------| | `num_experts` | `experts.num_experts`, `gate.weight.shape[0]`, or config attrs | Yes | | `dim` | `config.hidden_size` | Yes | | `moe_intermediate_size` | `experts.gate_up_proj.shape[1] // 2` or config | Yes | | `top_k` | `moe_block.top_k`, `config.num_experts_per_tok` | Yes | | `score_func` | `sigmoid` if `e_score_correction_bias` exists, else `softmax` | Blocker if neither | | `route_norm` | `config.norm_topk_prob` or default `True` | No | | `route_scale` | `config.routed_scaling_factor` or `1.0` | No | | `num_expert_groups` | `config.n_group` or `None` | No | | `shared_expert_info` | Probes `shared_expert`/`shared_experts`/`shared_mlp` | No | **Issues to detect and fix in experiment code:** - Expert weight layout transposed (e.g. `(E, H, 2*I)` instead of `(E, 2*I, H)`) — fix `state_dict_adapter.py` to transpose during conversion, and `moe_replacement.py` to detect the layout - Experts stored as `nn.ModuleList` of individual MLPs instead of fused 3D tensors — this is a true blocker (needs HF upstream change), report to user and stop - Scoring function not `softmax`/`sigmoid` but detectable from the model code (e.g. `torch.sigmoid` call without `e_score_correction_bias`) — fix `_resolve_score_func()` in `moe_replacement.py` - Router uses `nn.Linear` directly (key is `router.weight` not `gate.weight`) — add pattern to `state_dict_adapter.py` - MoE layer attribute is not `mlp` (e.g. `feed_forward`) — fix `model.py` and `moe_replacement.py` to handle the alternative name - Shared expert round-trip fails (singular/plural mismatch) — fix reverse patterns in `state_dict_adapter.py` - Model doesn't support `grouped_mm` expert implementation — adjust the numerical test to skip the `grouped_mm` alignment for this model - Layer-level MoE: router and experts are siblings of the dense MLP at the decoder layer level, not nested inside `layer.mlp` (e.g. Gemma4 has `layer.mlp` + `layer.router` + `layer.experts`). In this case the dense MLP is functionally a **shared expert** — it runs on every token and its output is summed with the routed expert output. The correct mapping is to treat the dense MLP as `shared_experts` inside Titan's `MoE` module and replace `layer.mlp` with the full Titan MoE (router + experts + shared_experts). This way the entire block has a single `ShardingConfig` and TP/EP work correctly. Fix in `moe_replacement.py`: when probing, detect that the dense MLP is a shared expert, include it in `shared_expert_info`, and when swapping, replace `layer.mlp` with the Titan MoE that includes the dense MLP as its shared expert. Wire the router and expert weights from the layer-level attributes into the Titan MoE. ### 1e. Probe attention for ShardingConfig Check which sharding configs `_set_layer_sharding_configs` (in `hf_sharding.py`) needs: Q and KV are detected **independently**, so any mix of full-rank/low-rank works (e.g. DeepSeek-V2-Lite has full-rank `q_proj` + low-rank KV): | Pattern | Projections | Already supported? | |---------|-------------|--------------------| | Full-rank Q | `q_proj` | Yes (`colwise_sharding_config`) | | Low-rank Q (MLA) | `q_a_proj`, `q_b_proj` | Yes (`_replicate_sharding_config`/`colwise_sharding_config`) | | Standard KV (GQA) | `k_proj`, `v_proj` | Yes (`colwise_sharding_config`) | | Low-rank KV (MLA) | `kv_a_proj_with_mqa`, `kv_b_proj` | Yes (`_replicate_sharding_config`/`colwise_sharding_config`) | | Output | `o_proj` or `dense` | Yes (`rowwise_sharding_config`) | | Q/K norms | `q_norm`, `k_norm` | Yes (`_replicate_sharding_config`, each independent) | | V-norm | `v_norm` | Yes (`_replicate_sharding_config`) | | DSA indexer | `indexer` | FSDP/EP only — **NOT under TP** (fails loud; the indexer's scatter_/index ops need the `spmd_types` backend, see `hf_sharding.py`) | The `_assert_all_states_sharded` backstop will fail loud at setup if any param/buffer-bearing attention (or other dense-path) module is left without a config, so an unhandled projection surfaces immediately. If found, report: "attention has unsupported projection `<name>`, needs `_set_layer_sharding_configs` update in `hf_sharding.py`." Also check: does the layer have `post_attention_layernorm`? (Some models omit it — already handled.) Does the embedding have extra buffers (e.g. Gemma4's `embed_scale`)? (Already handled — buffers are enumerated dynamically.) ### 1f. State dict adapter round-trip ```python from torchtitan.experiments.transformers_modeling_backend.state_dict_adapter import ( hf_to_titan_moe_state_dict, titan_to_hf_moe_state_dict, ) hf_sd = model.state_dict() titan_sd = hf_to_titan_moe_state_dict(hf_sd) roundtrip_sd = titan_to_hf_moe_state_dict(titan_sd) ``` Verify: all original keys present, no extra keys, values match. If keys are missing, the adapter needs new patterns for this model's key naming convention. ### 1g. Report Print a compatibility summary table: ``` Compatibility report for <model_id> model_type: <type> hidden_size: <H> moe_intermediate: num_experts: <E> top_k: <K> score_func: softmax|sigmoid route_norm: True|False shared_experts: none|additive|sigmoid-gated attention: gqa|mla|mla+dsa MoE layers: <indices> min_layers_needed: <N> state_dict_roundtrip: PASS|FAIL issues: none | <list> ``` ### 1h. Fix issues If issues were found, fix them in the experiment code now. Consult the "Issue Resolution Reference" at the bottom of this document for where each fix goes. The only true blocker is `nn.ModuleList` experts — for everything else, implement the fix and re-run the probe to verify. After all issues are resolved and the probe passes cleanly, proceed to Phase 2. ## Phase 2: Add Synthetic Test Config Add an entry to `_MODEL_CONFIGS` in `torchtitan/experiments/transformers_modeling_backend/.claude/skills/add_moe_model/scripts/numerical_equivalence.py`. ### Config rules Use production-scale `hidden_size` and `moe_intermediate_size` from the real config. Override everything else for a small single-GPU test: ```python "<model_type>": dict( model_type="<model_type>", hidden_size=<real>, # from pretrained config intermediate_size=<real>, # from pretrained config moe_intermediate_size=<real>, # from pretrained config num_local_experts=8, # reduced from real num_experts_per_tok=2, # reduced from real num_hidden_layers=<min>, # minimum for MoE layer to exist num_attention_heads=<fits>, # must divide hidden_size evenly num_key_value_heads=<fits>, # must divide num_attention_heads evenly vocab_size=256, max_position_embeddings=64, attn_implementation="sdpa", use_cache=False, # Include model-specific fields: # first_k_dense_replace, decoder_sparse_step, mlp_layer_types, # n_group, topk_group, n_shared_experts, norm_topk_prob, # routed_scaling_factor, q_lora_rank, kv_lora_rank, etc. ) ``` **Attention dimension rules:** - `num_attention_heads` must divide `hidden_size` evenly - `num_key_value_heads` must divide `num_attention_heads` evenly - For MLA models: use small `q_lora_rank` (16), `kv_lora_rank` (16), `qk_rope_head_dim` (8), `qk_nope_head_dim` (8), `v_head_dim` (16) - For DSA models: use small `index_n_heads` (2), `index_head_dim` (8), `index_topk` (8) **MoE layer selection rules:** - If `first_k_dense_replace > 0`: set it to 0 (force first layer to be MoE) - If model uses `decoder_sparse_step`: set to 1 - If model uses `mlp_layer_types`: set to `["sparse"]` **Expert group routing:** - If `n_group` is set, `num_local_experts` must be divisible by `n_group` and `num_local_experts / n_group >= 2` ## Phase 3: Run Numerical Equivalence **HF is the gold standard.** The numerical equivalence test compares Titan's MoE output against HF's output. The test must run HF's forward pass unmodified — do NOT patch, monkey-patch, or alter HF's computation to make the test pass. The only permitted modifications to close accuracy gaps are **temporary changes to titan core code**, which must be reverted after verification. Specifically: - Do NOT patch the HF model's activation function to match titan's - Do NOT replace HF's routing with titan's routing - Do NOT modify the test's comparison logic or thresholds - Do NOT skip or disable parts of the comparison - Do NOT cast HF tensors to different dtypes to hide precision diffs (the existing float32 gate cast is an exception — it aligns titan's autocast behavior, not the other way around) - If titan produces different results from HF, the gap analysis (Phase 3c) must identify WHY by temporarily changing titan to match HF, not by changing HF to match titan **Do NOT modify the pass/warn/fail thresholds.** The thresholds (KL < 1e-6 PASS, KL < 1e-3 WARN, KL >= 1e-3 FAIL) are calibrated for the known dispatcher precision gap. If a new model hits WARN, verify via Phase 3c that the diff is from known causes. If it hits FAIL, fix the integration, do not loosen the thresholds. ### 3a. Synthetic test ```bash CUDA_VISIBLE_DEVICES=<gpu> python \ torchtitan/experiments/transformers_modeling_backend/.claude/skills/add_moe_model/scripts/numerical_equivalence.py \ --models <model_type> ``` **Expected:** PASS with KL < 1e-6 and round-trip max_diff = 0.00. If it fails, diagnose: - `ValueError` from `_probe_hf_moe_block` — missing config attribute, fix the synthetic config - Shape mismatch in `load_state_dict` — state dict adapter needs update - `RuntimeError` from `grouped_mm` — check expert weight shapes - Numerical FAIL — investigate routing or weight transfer issue ### 3b. Pretrained test (if weights available) ```bash CUDA_VISIBLE_DEVICES=<gpu> python \ torchtitan/experiments/transformers_modeling_backend/.claude/skills/add_moe_model/scripts/numerical_equivalence_pretrained.py \ --model_dir <path> ``` **Expected:** PASS with KL < 1e-3 (known dispatcher precision gap). ### 3c. Verify numerical gaps via temporary core alignment (MANDATORY) **This step is NOT optional.** Execute it fully and automatically whenever any test shows WARN or higher. Do NOT ask the user before proceeding — apply patches, run tests, record results, report, revert. The entire sequence runs without user interaction. The goal: bring max_diff to 0.00 by temporarily making titan's implementation match HF's exactly. Every remaining diff must be experimentally explained — never claim a root cause without running code that proves it. **Rules:** - **Never give up on a patch because it crashes.** If a patch causes an error (e.g. dtype mismatch in scatter), fix the error by extending the patch. The patch is temporary — make it work, don't declare it "not applicable." - **Never claim a root cause without verification.** If you think "the diff is from bmm vs grouped_mm", prove it: make both sides use the same kernel and show the diff drops. If you can't change one side (e.g. HF rejects grouped_mm), change the other side (e.g. make titan use bmm temporarily, or manually compute the expert forward with the same kernel as HF using titan's weights). - **Adapt patches to the model's code path.** The known patches target specific code paths (e.g. f32 accumulation in `LocalTokenDispatcher.combine()`). If the model uses a different path, write an equivalent patch for that path. The principle is the same -- align titan's operations with HF's. - **Keep going until max_diff=0.00.** If the known patches don't get there, trace the forward step by step to find additional differences. **Procedure:** 1. **Record baseline** — note current max_diff and KL before any patches. 2. **Apply patches incrementally** — apply each applicable patch one at a time. After each, re-run the test and record the new max_diff. If a patch crashes, fix the crash (extend the patch as needed) and re-run. If a patch has no effect (e.g. top_k=1 makes reshape+sum trivial), record "no effect (top_k=1)" and move on. 3. **Handle kernel mismatches** — if the model doesn't support `grouped_mm` on the HF side, the test can't force the same kernel. In that case, verify experimentally: - Extract titan's expert weights after weight transfer - Run the expert forward manually using HF's kernel (e.g. `torch.bmm`) with titan's weights and the same input - Compare that output against HF's expert output - If they match (0.00), the remaining diff is confirmed as kernel difference, not a bug - Record this in the report as a verified kernel difference 4. **Investigate remaining diffs** — if max_diff is still non-zero after all patches and kernel alignment, trace the forward path step by step: router logits → routing weights → per-expert gate/up/down → score application → combine. Compare intermediate values between HF and titan to find where they first diverge. Make additional minimal patches to test hypotheses — if a patch doesn't help, revert it and keep digging. 5. **Report** — present a table to the user showing each patch and its measured impact: ``` Numerical gap analysis for <model_name> | Patch | max_diff | KL | Description | |-------|----------|----|-------------| | baseline | X.XXe-XX | X.XXe-XX | no patches | | +patch 1 | X.XXe-XX | X.XXe-XX | f32 accumulation in combine | | +patch 2 | X.XXe-XX | X.XXe-XX | reshape+sum (no effect, top_k=1) | | +kernel align | X.XXe-XX | X.XXe-XX | titan experts via bmm | | final | 0.00e+00 | X.XXe-XX | all diffs explained | ``` 6. **Revert ALL patches** — restore all core files: ```bash git checkout torchtitan/models/common/moe/token_dispatcher.py torchtitan/models/common/moe/moe.py ``` Verify with `git diff` that the revert is clean. The temporary patches are NOT part of the commit. **Known sources of numerical difference:** Do NOT copy-paste patches from this document — the titan core code may have changed since these were written. Instead, read the current code in `torchtitan/models/common/moe/token_dispatcher.py` and `torchtitan/models/common/moe/moe.py`, understand the principle behind each fix, and write a patch that applies cleanly to the current code. The three known differences and how to fix them: **1. bf16 accumulation before scatter_add** — `LocalTokenDispatcher.combine()` multiplies expert output by routing scores in f32 but casts back to bf16 before `deterministic_scatter_add`, so the accumulation happens in bf16. **Fix principle:** remove the cast to bf16, keep everything in f32 through the scatter_add, allocate an f32 output buffer, cast to the input dtype only at the very end. Example (illustrative, may not match current code): ```python # principle: remove .to(routed_output.dtype), use f32 buffer routed_output = routed_output.to(torch.float32) * scores out = scatter_add(f32_buffer, ..., routed_output) return out.to(x.dtype) ``` **2. scatter_add vs reshape+sum accumulation order** — `scatter_add` accumulates top_k contributions in expert-sorted order. HF uses `reshape(N, K, D).sum(dim=1)` which accumulates in token order (consecutive top_k entries per token). Different f32 summation order produces different results. **Fix principle:** store the argsort permutation from dispatch, in combine unsort back to original token order, then reshape and sum over the top_k dimension. For `top_k=1` this has no effect (only one element to sum). **3. topk sorted flag** — titan uses `torch.topk(sorted=False)`, HF defaults to `sorted=True`. With `sorted=True` topk values come back in descending order; the normalization sum is computed over differently ordered values, producing a small f32 rounding difference. **Fix principle:** change `sorted=False` to `sorted=True` in the router's topk call. For `top_k=1` this has no effect (only one element). With all three fixed, **all existing models produce max_diff=0.00.** If the new model still shows non-zero diff after all applicable fixes, the remaining difference is model-specific. Investigate as described in step 4 above. ## Phase 4: Parallelism Integration Tests Run the new model through the same parallelism configurations used for existing models. The integration test infrastructure is in `torchtitan/experiments/transformers_modeling_backend/tests/integration_tests.py`. ### 4a. Determine available GPUs ```bash nvidia-smi -L | wc -l ``` ### 4b. Run parallelism configs Test the new model with at least these configurations (adapt GPU counts to what is available). Add three temporary Python recipes to `config_registry.py`; configuration flags are not accepted by the launcher: ```python def _new_model_test_config(*, tp: int = 1, ep: int = 1) -> Trainer.Config: config = transformers_modeling_backend_debugmodel_moe() model_config = model_registry( "debugmodel_moe", seq_len=config.training.max_context_length, hf_model="<model_id>", ) config.model = model_config config.loss = CrossEntropyLoss.Config(global_vocab_size=model_config.vocab_size) config.training.steps = 2 config.parallelism.data_parallel_shard_degree = -1 config.parallelism.tensor_parallel_degree = tp config.parallelism.expert_parallel_degree = ep return config def new_model_fsdp() -> Trainer.Config: return _new_model_test_config() def new_model_fsdp_ep() -> Trainer.Config: return _new_model_test_config(ep=<EP>) def new_model_fsdp_tp_ep() -> Trainer.Config: return _new_model_test_config(tp=2, ep=<EP>) ``` Set `config.hf_assets_path` in the helper as well if the test needs the new model's tokenizer assets. Keep `training.steps = 2`: the goal is to verify the model initializes, parallelizes, and runs forward/backward without errors, not to converge. ```bash # FSDP only (minimum viable test) NGPU=<NGPU> MODULE=transformers_modeling_backend CONFIG=new_model_fsdp \ ./run_train.sh # FSDP + EP NGPU=<NGPU> MODULE=transformers_modeling_backend CONFIG=new_model_fsdp_ep \ ./run_train.sh # FSDP + TP + EP (if NGPU >= 4) NGPU=<NGPU> MODULE=transformers_modeling_backend CONFIG=new_model_fsdp_tp_ep \ ./run_train.sh ``` **Every config must be run.** There are only three — run all of them. Do NOT skip, do NOT mark as "NOT TESTED", do NOT mark as "SKIPPED". Every config in the final report must show either PASS or FAIL with a diagnosed error message and root cause. If a config fails: 1. Read the error traceback 2. Diagnose the root cause 3. If fixable in experiment code, fix it and re-run 4. If it's a fundamental limitation (e.g. DTensor mixing), report FAIL with the specific error and why it can't be fixed in experiment code **Redundant compute is a blocker.** If the only way to make TP work for a module is to run the full computation redundantly on every TP rank (e.g. all-gather input, run full forward, slice output), that is NOT an acceptable fix — it defeats the purpose of TP. If a model component (e.g. a custom attention mechanism) cannot be properly sharded, report it as a FAIL with the reason "no proper TP sharding available, would require redundant compute". Do NOT implement redundant-compute workarounds with pre/post hooks that all-gather and re-shard. **Use ShardingConfig, not hooks.** The HF backend sets ``_sharding_config`` on every sub-module so ``model.parallelize()`` handles all DTensor distribution and forward wrapping. Available helpers from ``torchtitan/models/common/decoder_sharding.py`` and ``torchtitan/experiments/transformers_modeling_backend/hf_sharding.py``: - ``colwise_sharding_config()`` — for linear projections (weight Shard(0), output Shard(-1)) - ``rowwise_sharding_config(output_sp=True)`` — for output projections (weight Shard(1), output Shard(1) for SP) - ``_hf_norm_sharding_config(enable_sp=True)`` — for norms (weight Replicate, activations Shard(1) when SP enabled) - ``_replicate_sharding_config(module)`` — for modules that should have Replicate params/buffers without TP sharding (dynamically enumerates params and buffers to avoid ``_shard_states`` errors) - ``ShardingConfig(in_src_shardings=..., in_dst_shardings=...)`` — for input redistribution at module boundaries (e.g. attention, MLP gather from Shard(1) to Replicate) Do NOT use ``register_forward_pre_hook`` / ``register_forward_hook`` to manually convert DTensor ↔ local tensors. Do NOT use ``parallelize_module`` from ``torch.distributed.tensor.parallel``. Use ``_sharding_config`` exclusively. Pass `<model_id>` to `model_registry` in the temporary recipes. If the base config has incompatible defaults for the new model (e.g. wrong tokenizer path), update those config fields in the helper as well. Do NOT give up because "the config is incompatible" -- fix the recipe. Common failure modes and fixes: - TP failures from unsupported attention projections — fix in `hf_sharding.py` by adding ShardingConfig in `_set_layer_sharding_configs` - EP failures from expert weight sharding — check FSDP `shard_placement_fn` - Shape mismatches from model-specific layer structure — fix in experiment code - Config incompatibility -- update the temporary recipe - Tokenizer mismatch -- set `config.hf_assets_path` to the model assets - Custom attention with no standard Q/K/V/O projections — FAIL, not fixable without redundant compute - Mixed DTensor / plain Tensor — a module has buffers not declared in its ShardingConfig's ``state_shardings``. Fix by using ``_replicate_sharding_config(module)`` which dynamically enumerates all params and buffers **Note:** If no GPUs are available, skip this phase and note it in the report. But if GPUs are available, all configs must be run and reported. ## Phase 5: Update Model Compatibility Doc Update `torchtitan/experiments/transformers_modeling_backend/MODEL_COMPATIBILITY.md` with the new model: 1. Add a row to the **Numerical Equivalence Summary** table with the model's status, KL, max_diff, cos_sim, and round-trip result. 2. Add a section under **Architectural Differences per Model** describing: - Router type and features (softmax/sigmoid, normalization, bias, scaling) - Expert format (standard/transposed, activation function) - Shared experts (none/additive/sigmoid-gated/layer-level) - Attention type (GQA/MLA/DSA/custom) - Any unsupported features and their impact - Summary of differences from Titan's MoE 3. Add a row to the **Parallelism Support** table (FSDP and FSDP+EP). 4. If the model has unsupported features requiring core changes, add or update entries under **Core Changes Needed for Full Support**. ## Phase 6: Lint and Commit 1. Run pre-commit on all changed files: ```bash SKIP=pyrefly-check pre-commit run --files <changed_files> ``` Fix any failures and re-run until clean. 2. Verify no core files were modified: ```bash git diff --stat # only files in experiments/transformers_modeling_backend/ ``` 3. Create commit: ``` Add <ModelName> to transformers MoE backend Adds synthetic test config with production-scale dimensions (H=<hidden_size>, I=<moe_intermediate_size>, E=8). Verified via numerical equivalence test (KL=<value>, max_diff=<value>). ``` ## Issue Resolution Reference Most issues are fixable in the experiment folder. Only `nn.ModuleList` experts is a true blocker requiring upstream changes. | Symptom | Where to fix | What to do | |---------|-------------|------------| | Expert weights transposed `(E, H, 2*I)` | `state_dict_adapter.py` | Add transpose in conversion; fix shape detection in `moe_replacement.py` | | `experts` is `nn.ModuleList` | **Cannot fix** — needs HF upstream | Report to user, stop | | Sigmoid scoring not auto-detected | `moe_replacement.py` | Add detection path in `_resolve_score_func()` (e.g. check router source for `sigmoid`) | | Router key `router.weight` not converted | `state_dict_adapter.py` | Add regex pattern in `_build_hf_to_titan_patterns()` | | Shared expert round-trip broken | `state_dict_adapter.py` | Fix reverse patterns in `_build_titan_to_hf_patterns()` | | MoE layer attr is `feed_forward` not `mlp` | `model.py`, `moe_replacement.py` | Handle alternative attribute name | | Attention projection not in ShardingConfig | `hf_sharding.py` | Add case to `_set_layer_sharding_configs` using `colwise_sharding_config()`/`_replicate_sharding_config()` | | Model rejects `grouped_mm` | `scripts/numerical_equivalence.py` | Skip `_experts_implementation` for this model | | `trust_remote_code` import error | Config loading | Try without `trust_remote_code` first | | `num_hidden_layers=1` has no MoE layer | Test config | Set `first_k_dense_replace=0` or increase layers | | State dict keys missing after round-trip | `state_dict_adapter.py` | Add patterns in `_build_hf_to_titan_patterns()` / `_build_titan_to_hf_patterns()` | | Layer-level MoE (router/experts are siblings of dense MLP) | `moe_replacement.py` | Treat dense MLP as shared expert inside the Titan MoE; replace `layer.mlp` with full MoE block containing router + experts + shared_experts(=dense MLP) | | TP fails with mixed DTensor/Tensor | `hf_sharding.py` | Module has undeclared buffers — use `_replicate_sharding_config(module)` which dynamically enumerates all params and buffers | | Embedding has extra buffers (e.g. `embed_scale`) | `hf_sharding.py` | Already handled — `set_hf_sharding_configs` enumerates embedding buffers dynamically | ## Files touched (experiment folder only) | File | Change | |------|--------| | `scripts/numerical_equivalence.py` | Add entry to `_MODEL_CONFIGS` | | `state_dict_adapter.py` | Add key patterns (only if round-trip fails) | | `hf_sharding.py` | Add ShardingConfig entries (only if new projection/norm names) | | `moe_replacement.py` | Add probing logic (only if new MoE pattern) |  --- name: inference_perf_hillclimb description: Profiler-driven hill-climbing to close the inference throughput gap between TorchTitan's unified model (running inside vLLM) and vLLM's native model. Benchmark with generate.py --benchmark, climb optimization rungs (compile / cudagraph / fused kernels), profile torchtitan vs the native target, then patch the single biggest gap at a time and re-measure. Use when the user wants to benchmark or optimize RL inference generation speed, reproduce previous hill climbing study, or invokes /inference_perf_hillclimb. --- # Inference Performance Hill-Climbing Close the gap between the TorchTitan unified model (in vLLM) and vLLM's native model on inference throughput, one profiler-found optimization at a time. Harness: `torchtitan/rl/generate.py --benchmark`. Append-only running log + full per-rung numbers live in `torchtitan/rl/docs/inference_gap_ablation.md`. Goal: match vLLM native 100%. ## Scope: close the IMPLEMENTATION gap, not the serving knobs This skill makes torchtitan's model run as fast as vLLM's NATIVE model at a FIXED model, workload, precision, and topology -- i.e. it closes the framework/overhead gap (target = native at the *same* config; ratio -> 1.0). It is explicitly NOT about the orthogonal levers that change the setup itself: - increasing batch size / seq_len / amount of input data - changing the model - lower precision (fp8 / quantization) - speculative decoding - disaggregated prefill - changing topology / number of GPUs Those move ABSOLUTE throughput but apply equally to native -- they don't change the torchtitan-vs-native ratio this skill chases. Hold them FIXED and IDENTICAL between torchtitan and native when measuring; otherwise you're no longer measuring the implementation gap. ## Before any run: confirm with the user (don't assume defaults) - **Model + parallelism**: which model, TP degree, SP/DP/PP/EP (e.g. Qwen3-32B, TP=8, SP off). Sets the sharding, all-reduce sizes, cudagraph capture sizes, and what "native" even is -- the dominant gap and right knobs change with it. - **Versions**: vLLM / torchtitan / PyTorch-nightly. They move fast and change kernels, cudagraph behavior, custom AR, and APIs. Record all three -- numbers across versions are not comparable. - **Workload**: input_len, batch_size, gen_tokens. Standard: W1 bs8/in1024/out128 (priority); W2 bs32/in4096/out1024 (needs `--max-seq-len 8192`). Sets prefill-vs- decode weighting, AR/capture sizes, memory -- the dominant gap depends on it (e.g. prefill capture helps short-gen W1 more than long-gen W2). - **compile + cudagraph** (the two settings that dominate perf; use AskUserQuestion): - compile: `aot_eager` (validated torchtitan default, per-layer torch.compile) | `inductor` (heavier) | `off`. **The fair target is native(eager)**, NOT native(inductor): inductor fuses the TP all-reduce into gemm/RMSNorm (`fuse_allreduce_rms`), eager does not -- it's ~+5% and not in the aot_eager path, so comparing to it invents a fake gap. - cudagraph mode: pick ONE for baseline AND target; never mix across rungs (see "cudagraph modes" under Harness). ## The hill-climbing loop (the core method -- PROFILE-DRIVEN, measure don't guess) **Every rung is justified by a PROFILE, never intuition**, and KEPT only if a re-benchmark moves throughput. The flag list is the residue of past profiles; the method is: profile -> single biggest gap -> patch exactly that -> re-measure -> re-profile. Two rules: (a) **trace time misleads** -- a huge all-reduce duration is usually cross-rank SPIN (sync wait), not reducible compute; confirm with tok/s. (b) **the biggest gap is often NOT a kernel** -- prefill cudagraph capture and arrival-spin were the top levers; profile broadly. 1. Confirm model / parallelism / versions / workload / compile+cudagraph (above). 2. Establish the ladder once, capture the **target**: `--native --compile aot_eager --cudagraph on` (same workload). 3. **PROFILE both** torchtitan-best and native (`--profile`). Decompose for the gap: - kernels: `/tmp/ablation_logs/analyze_trace.py` (single-rank diff: kernel us, extra kernels present in one but not the other, CPU-op deltas). - collectives: `/tmp/ablation_logs/analyze_ar_spread.py` over the 8 per-rank traces -- min(dur)=pure comm, spread(dur)=ARRIVAL SPIN. Single-rank kernel sums mislead for collectives. - host-bound: per-rank GPU busy% + idle-gap location (eager prefill vs captured decode). Low busy% on the driver rank => host-launch-bound. 4. Pick the SINGLE biggest gap (kernel / spin / host lag / capture coverage). 5. Implement exactly that as a new knob (custom_op patch in `kernel_ablation.py`, a `--model-2d` variant, or a cudagraph/config flag); often "make torchtitan do what native does here". 6. **Validate numerics**: example mode (no --benchmark) at TP>=2, greedy text bitwise-identical with vs without the patch. 7. Commit the rung (one git commit per table line). 8. Re-benchmark (W1, then W2); append to the doc. KEEP only if throughput moved (trace promise != throughput). 9. Repeat from step 3 (RE-PROFILE -- the biggest gap shifts after each patch) until within ~3% of target, or the remainder is structural (document it). ## Gap buckets: how to spot, what we found, the lever - **CPU / host overhead** (DTensor dispatch, ATen dispatch, per-launch Python). In an EAGER region the host can't feed the GPU -> it starves. Spot: per-rank GPU busy% < ~90% with idle gaps before compute kernels. Levers: cudagraph capture takes the CPU out of the per-kernel path (capturing PREFILL too was the single biggest lever, ~+20%); OR remove DTensor from the forward (`--model-2d` pure-local); OR cut per-launch overhead. Hidden under a captured graph -- only bites in eager prefill. NOTE: the structural lever is removing DTensor, not "2D" -- compile+cudagraph kill DTensor's CPU cost but its boundaries (from_local / subclass dispatch / placement views) fragment the captured graph -> launch jitter. Tensor rank is a red herring (3D==2D); register_sharding REGRESSES (-26%, more DTensor dispatch). - **Single kernel slowness** (same shape, torchtitan > native). Spot: `analyze_trace.py` per-kernel us. Found: NCCL all-reduce 23us vs vLLM custom one-shot AR 6.3us for small decode messages (a real algorithm difference); faster RoPE/attn variants. Lever: `--allreduce-vllm`, `--rope-kernel helion`, `--attn-backend flash`. Caveat: GEMM and RMSNorm are already at native parity (torch/Quack `rms_norm` even BEATS vLLM's) -- don't chase them. - **Fused vs separate kernels** (native fuses what torchtitan runs as N kernels: fused add+RMSNorm, fused QKV / gate-up GEMM, SiluAndMul). Spot: "extra kernels" in the torchtitan trace + higher CPU-op count. Lever: `fuse_qkv=True` + `fused_swiglu`, `--fused-addnorm`, `--silu-vllm`. Caveat: mostly a launch-count / host win -> matters in eager regions, shrinks under cudagraph. - **Communication straggler** (one rank reaches the collective late, the other 7 SPIN). Spot: `analyze_ar_spread.py` (min(dur)=comm==native, spread=arrival spin). Find WHAT BOUNDS the straggler: *compute?* per-rank compute time is usually UNIFORM, so no; *host-launch?* the DRIVER rank (rank 0) runs scheduler/sample/ output on its Python thread and falls behind in the eager prefill -> arrives last (busy 54% vs 94%, spin ~600us) -- the usual cause; fix = capture prefill (spin ~600us -> ~3us, the +20% lever) or lighten the per-step host path. *hardware/topology?* if the SAME rank straggles in NATIVE too, it's a GPU/NVLink-position effect, not fixable in the model. Rule: the AR kernel is innocent -- fix uneven ARRIVAL, not the collective. Operational gotchas: check GPU contention first (`nvidia-smi`; a co-tenant job spikes variance to +/-150 vs clean +/-1-7 -- re-run clean); NO Python-side logging/prints inside a patched forward (forces a torch.compile graph break); launch long runs detached (`setsid nohup ... &`) and clean GPU stragglers between. ## Harness **Build/restore**: stock `generate.py` on main is an EXAMPLE (single prompt). Bring the benchmark harness over from branch `ablation-inference`: `git checkout ablation-inference -- torchtitan/rl/generate.py` + the patch modules (`models/{kernel_ablation,qwen3_vllm_2d,helion_rope,vllm_fused_ops}.py`), fix the import (`from torchtitan.rl.examples.alphabet_sort import config_registry`), and add `rl_grpo_qwen3_32b` (mirror `rl_grpo_qwen3_14b`, `model_registry("32B")`, TP=8). `--benchmark` builds the engine OUTSIDE the timed region (excludes startup/compile/capture), runs `--warmup-runs` then times `--num-runs`, reports tok/s = batch*gen/wall; feeds synthetic token-ids (skips the tokenizer); sets `enable_prefix_caching=False`. Build the cudagraph `CompilationConfig` EXPLICITLY (the stock helper only emits `cudagraph_mode="full"`). **Launch**: `torchrun --nproc_per_node=<TP> generate.py --benchmark ...`. Env: conda `titan-rl` (see `torchtitan/rl/README.md` to build it). **Knobs** are the residue of past profiles -- NOT a required checklist, and NOT guaranteed to exist on a fresh checkout (the harness is rebuilt each time; re-add what you need, add NEW knobs for new gaps): - **Kernel patches** (torchtitan model; stack as the profile dictates): `--rope-kernel {helion,vllm} --silu-vllm --rmsnorm-vllm --allreduce-vllm --fused-addnorm --attn-backend {custom,flash}` -- custom_op monkeypatches in `models/kernel_ablation.py` / `vllm_fused_ops.py` (survive compile+cudagraph). FusedQKV+gate_up: use config `fuse_qkv=True` + `overrides/fused_swiglu.py` `@override`, not the superseded `--merged-gemm`. - **Whole-model paths** (mutually exclusive with kernel patches): `--model-2d {local,localfused,spmd,local3d,dtensor}`. `localfused` = pure-local (no DTensor) + fused add+RMSNorm (best); `spmd` = localfused but the TP AR routes through `spmd_types.redistribute(P->R)`. spmd_types is a pre-run CHECK (validates SPMD sharding via typechecking) -- keep it OFF for perf runs. - **compile / cudagraph**: `--compile {off,aot_eager,inductor}` `--cudagraph {on,off}` `--cudagraph-mode {full_decode_only,full}`. Capture PREFILL with `--cudagraph-mode full --max-num-batched-tokens --max-capture-size `, where P >= the prefill CHUNK size (= max_num_batched_tokens), NOT input_len. Decode capture sizes default to powers of 2 up to max_num_seqs (= batch). - **misc**: `--nccl-algo`, `--native` (target), `--profile` (per-rank chrome trace), `--max-seq-len` (long prompts). ### cudagraph modes (pick ONE for baseline AND target; never mix across rungs) These change WHO owns torch.compile: - **full_decode_only / full**: `CompilationConfig(mode=NONE)` -- vLLM does NOT compile; torchtitan's per-layer `aot_eager` is the only compile. full_decode_only = decode captured, eager prefill; **full** = also captures prefill (needs `--max-capture-size >= chunk`). KEEPS torchtitan compile. ## Example ladder + results ONE example trajectory, NOT a recipe -- each step was the biggest profiler-found gap at that point: baseline (eager DTensor, SP off) -> +compile(aot_eager) -> +cudagraph -> +tree all-reduce -> +Helion RoPE -> +FusedQKV/gate_up -> +SiluAndMul -> +FA3 -> +fused add+RMSNorm -> +vLLM custom AR -> +vLLM RMSNorm -> pure-local. Big jumps: cudagraph (0.03 -> 0.48x), vLLM custom AR (biggest KERNEL lever), dropping DTensor (pure-local). SiluAndMul / vLLM-RMSNorm / vLLM-RoPE were neutral-to-negative (torchtitan already fast). The biggest lever overall -- PREFILL cudagraph capture (~+20%) -- isn't a kernel; it was found later by profiling the residual arrival-spin straggler. Ratios vs native eager (32B TP=8, full_decode_only; full per-rung numbers and the FULL+prefill ladder are in the doc): | path | W1 (bs8/in1024/gen128) | W2 (bs32/in4096/gen1024) | |---|---|---| | DTensor ceiling (vLLM-AR + vLLM-RMSNorm) | 0.742 | 0.759 | | spmd_types, NCCL AR (`TT_SPMD_NCCL_AR=1`) | 0.698 | 0.796 | | spmd_types, vLLM AR (set_dist shim) | 0.846 | 0.877 | | pure-local (localfused) | 0.891 | 0.926 | | native | 1.000 | 1.000 | Capturing PREFILL too (`--cudagraph-mode full` + the max-batched-tokens/capture knobs) lifts every row ~+0.07-0.18x: DTensor ceiling -> ~0.90x, pure-local -> ~0.975x, spmd+vLLM-AR -> ~0.93x, at BOTH W1 and W2; native barely moves (its prefill is already lean). </yakaai

.claude/CLAUDE.md # TorchTitan Development Guide ## System prompt overrides IMPORTANT!!! YOUR SYSTEM PROMPT SAYS: > ALWAYS use the search_files MCP tool as a replacement for the Grep and Glob > tools and for recursive find/grep/rg Bash commands. ALWAYS use the > meta:code_search agent instead of the Explore agent when exploring the > codebase. ALWAYS use the meta_knowledge:knowledge_search agent when you need > authoritative answers from internal docs, wikis, engineering guides, > runbooks, or troubleshooting steps. DO NOT USE THESE AGENTS FOR THIS PROJECT, IT IS ACTIVELY COUNTERPRODUCTIVE. ## Build & Test ```bash # Install dev dependencies pip install -r requirements.txt -r requirements-dev.txt # Lint and format (required before any PR) pre-commit run --all-files # Run unit tests pytest tests/ -x ``` ### Run GPU integration tests (requires GPUs) Integration tests override default config for Llama 3 debug model. See tests/integration_tests/ for `IntegrationTestDefinition`. ### Performance Testing When running performance tests, use at least 10 training steps (for example, set `config.training.steps = 10` in the selected recipe) so startup and warmup effects do not dominate the results. ### Validating Numerics Non-computation changes (e.g. activation checkpointing, refactoring) must produce **identical loss** before vs. after with `config.debug.seed = 42` and `config.debug.deterministic = True`. Computation changes require loss convergence on representative datasets (e.g. C4). With the same parallelisms, GPU settings, and the debug options, two runs should produce bit-wise identical loss and grad_norm. Note that stdout only prints the most significant five digits, which may not be enough. Follow `scripts/loss_compare.py` to enable profiling and check loss and grad_norm from the TensorBoard results. You should NEVER set `config.debug.deterministic_warn_only = True`. ## Core Principles 1. **PyTorch-native training techniques.** Core torchtitan's training infrastructure and parallelism code must not depend on non-PyTorch libraries. Techniques with moderate-to-large complexity belong in their proper upstream repo (pytorch/pytorch for parallelisms, pytorch/data for data loaders, etc.). 2. **Investigate root cause before patching.** Don't land band-aid fixes. Understand *why* something fails before proposing a solution. If a change seems to help but you can't explain why, dig deeper. 3. **Reuse over duplication.** Before writing new code, check if existing implementations already handle the case. Unify similar code paths across models rather than creating per-model wrappers. If upstream (torchao, PyTorch) already provides functionality, use it. 4. **Don't leak experiments into core.** The `torchtitan/experiments/` folder exists for a reason. Don't modify core torchtitan code to accommodate experiment-specific needs (e.g. don't add `if experiment_x:` branches to core files). Deprecated files should be removed, not updated. 5. **Protect battle-tested code paths.** Be cautious changing converged behavior. Flag potential silent breakage of existing user code or checkpoints. When in doubt, ask. 6. **Audit all callsites.** When changing shared code (common model components, config fields, distributed utilities), check and update every callsite. This includes all model variants: llama3, qwen3, qwen3_5, deepseek_v3, deepseek_v4, gpt_oss, kimi_k2_7, kimi_k3, muse_glimmer, flux, etc. 7. **No speculative defensive checks.** Don't add checks, casts, fallbacks, or conversions "just in case." Only validate explicit contracts, user-facing configuration, or invariants whose failure would otherwise be silent or unclear. ## Code Style ### Unicode ASCII only in newly added or rewritten code comments and docstrings. Don't introduce Unicode characters (e.g. smart quotes, em dashes, arrows, non-ASCII letters) in comments or docstrings you add or change. Use ASCII equivalents: `->` for arrows, `<-` for left arrows, `<->` for bidirectional arrows, `--` for em dashes. Leave preexisting Unicode in untouched comments alone; only enforce this for the comments and docstrings you are adding or rewriting. ### Naming - Names must be **accurate, descriptive, and reflect actual scope**. Don't use "toy/test/temp" in production names -- put that context in docstrings instead. - Follow upstream conventions: match torchao and PyTorch naming where applicable. E.g. if torchao calls it `Float8Linear`, use `Float8Linear` not `Float8Config`. - Use `num_` prefix for counts (e.g. `num_expert_groups` not `n_expert_groups`) when not directly matching an upstream API. - **`axis` names a specific mesh axis; `dim` is for tensors and mesh shape.** In any name we own -- variables, parameters, attributes, helpers, comments, docstrings, error messages -- use ``axis``/``axes`` when referring to a specific ``DeviceMesh`` axis (TP axis, ``dp_shard`` axis, the list of axes a spec references). Use ``dim``/``dimensional`` for the mesh's *shape* ("1D mesh", "multi-dimensional SPMD mesh") and for tensor dimensions; bare ``dim`` on its own should refer to a tensor dimension. The exception is when calling into PyTorch upstream APIs (``DeviceMesh.mesh_dim_names``, ``DataParallelMeshDims``, etc.): match the upstream spelling at the call site, then assign into a locally named ``mesh_axis_names`` if the value flows through our code. - **Shape-suffix tensor names.** In model code, name tensors with shape suffixes (Noam Shazeer convention: https://medium.com/@NoamShazeer/shape-suffixes-good-coding-style-f836e72e24fd), e.g. `x_BLD`, `q_BLHK`, `out_THV`. Capital-letter suffixes denote *logical* tensor dimensions, not a physical sharding layout -- a name like `routed_input_RD` keeps the same suffix whether or not `R` is a local shard under EP/SP. Letters are scoped per module, not global: give each module that uses them a legend in a top-of-file comment, and don't assume a letter means the same thing across files (e.g. `N` is num heads in `attention.py` but routed tokens in `moe.py`). Apply this to newly written or rewritten model tensor code; don't churn unrelated code just to add suffixes. ### Code Placement - Put code in the **most general applicable location**: - Model-agnostic parallelism helpers -> `torchtitan/distributed/` - Model base classes every model implements or the engine consumes (`BaseModel`, `Module`, `ShardingConfig`, `StateDictAdapter`) -> `torchtitan/models/base/` - Shared, concrete model components (attention, MoE, embeddings, `Decoder`) -> `torchtitan/models/common/` - Model-specific code -> the specific model folder - Don't put model-agnostic functionality in model-specific files just because that's where you first needed it. ### Assertions and Error Handling - **`ValueError`** for user-facing errors (bad config, invalid input). - **`assert`** only for internal invariants that indicate programmer error. - Always validate mesh axes, tensor placements, and config values explicitly in distributed code -- don't assume a 1D mesh or specific placements. - When a code path silently skips user configuration, **emit a warning**. ### Parameters and Config - Important parameters first; less important ones later. - Prefer keyword-only arguments after the first positional arg. - No `None` defaults for required config fields. - `dataclasses.replace()` is a shallow copy: nested dataclasses and list/dict fields are shared by reference. Be explicit when deep copies are needed. ### Comments and Documentation - Add comments only for genuinely non-obvious things: dimension semantics, parallelism gradient placements, why a workaround exists. - Use TODO comments for known limitations with a brief explanation. - Put descriptions in docstrings, not in names. ## PR Expectations 1. **Lint first.** Run `pre-commit run --all-files` and fix all issues before requesting review. CI linting failures waste everyone's time. 2. **Show numerical proof.** Include loss comparison for any non-trivial change. 3. **Explain "why" not just "what"** in the PR description. 4. **Add tests.** New features need CPU unit tests at minimum; GPU integration tests when involving parallelism. Verify CI actually runs the intended test recipe (check the selected config function and its model configuration). 5. **Keep model code minimal.** After model changes, ensure original checkpoints still load correctly. Document reasons for model changes. --- description: Rules for the config system globs: torchtitan/** --- # Config System Rules ## No None Defaults for Required Fields If a config field is required, don't give it a `None` default. Either set a meaningful default or make it a required argument with no default. ## Parameter Ordering - Put important, commonly-used fields first. - Less important or rarely-changed fields go later. - Prefer keyword-only arguments after the first positional arg to prevent positional mistakes. ## Naming Conventions - Use descriptive names that reflect what the config controls. - Use `num_` prefix for count fields. - Prefer strings as config names when they benefit extensibility (e.g. model selection, dataset selection). ## Safety - When config values are passed through to external APIs (e.g. DataLoader kwargs), ensure all values are safe for all code paths (e.g. what happens when `num_workers == 0`?). - When a config option silently doesn't take effect in certain code paths, emit a warning to the user. --- description: Rules for distributed training code globs: torchtitan/distributed/** --- # Distributed Code Rules ## Assert Mesh and Placements Explicitly - Never assume a 1D mesh. Always assert on mesh axes before using them. (``axis``/``axes`` names a specific ``DeviceMesh`` axis like TP or ``dp_shard``; ``dim``/``dimensional`` describes the mesh's shape, e.g. "1D mesh" or "multi-dimensional mesh". Bare ``dim`` is reserved for tensor dimensions.) - Validate tensor placements (Replicate, Shard, Partial) explicitly. - When enforcing a field is not None for plain tensor inputs, do so with a clear error message. ## Document Parallelism Semantics For any usage of `DTensor.to_local`, specify explicitly the `grad_placements`, especially when the original DTensor placement has `Replicate` or `Partial`. This includes the indirect calls from `local_map` (`in_grad_placements`), and `full_tensor` (`grad_placements`). ## Consider All Parallelism Combinations When adding or modifying distributed code, think through how it interacts with every parallelism dimension. A fix for one configuration may break another. Include the parallelism configuration in bug reports and test descriptions. ## Model-Agnostic Code Lives Here Helper functions that apply to any model belong in `torchtitan/distributed/`, not in model-specific infrastructure files. ## Be Conservative with Changes Distributed training code is hard to test exhaustively. When modifying existing behavior: - Verify numerics match before and after across multiple parallelism configs. - Watch for silent correctness issues (wrong gradient placements, identity operations that break DCP). - If changing something that's been converged and validated, provide strong justification and thorough testing. --- description: Rules for the experiments folder globs: torchtitan/experiments/** --- # Experiments Folder Rules ## Looser Standards, Still Linted Code in experiments has more flexibility than core, but must still pass `pre-commit run --all-files`. No exceptions for linting. ## Don't Modify Core for Experiments If an experiment needs behavior that core doesn't provide, work around it in the experiment folder. Don't change core torchtitan code to accommodate experiment needs (e.g. don't add `if experiment_x:` branches to `train.py` or core modules). The experiments folder is "more-or-less a hack" by design. ## Use TorchTitan's Config System Use torchtitan's existing config and job config infrastructure. Don't introduce custom argument parsing or parallel config systems. ## Separate Concerns Keep distinct features in separate folders (e.g. inference code separate from RL work). Don't bundle unrelated functionality just because it's all "experimental." --- description: Rules for model implementations globs: torchtitan/models/** --- # Model Code Rules ## Keep Models Minimal and Readable - Model files should contain model architecture, not training infrastructure. The exception is a workload with its own training loop (Flux diffusion), which keeps its `Trainer`/`Validator` subclass, data and inference next to the model. - Weight initialization belongs in the config or a dedicated init function, not scattered through `Module.__init__`. - After any model change, ensure original checkpoints still load correctly. ## Audit All Variants When changing shared components (attention, normalization, MoE routing), check and update **all** model variants under `torchtitan/models/`: llama3, qwen3, qwen3_5 (and its qwen3_6/qwen3_8 flavors), deepseek_v3, deepseek_v4, gpt_oss, kimi_k2_7, kimi_k3, muse_glimmer, flux, and any others present. Don't leave stale patterns in sibling models. ## Unify Across Models - Don't create per-model wrappers for the same functionality. Aim for at most one general wrapper shared by all models. - If multiple models have near-identical code (e.g. sharding helpers or `_apply_fsdp` overrides), consolidate into `torchtitan/models/common/` or `torchtitan/distributed/`. - Before adding a new rotary embedding, MoE router, or similar component, check if an existing implementation already supports the use case. ## Standard Model Folder Structure Each model folder follows a consistent pattern (see `torchtitan/models/README.md`): - `flavors.py` -- `MODEL_FLAVORS` and `build_model_config(flavor, parallelism=...)` (sizes, hyperparameters, and each component's sharding) - `model.py` -- model definition; extra architecture files (e.g. `moe.py`) as needed - `sharding.py` -- model-specific `ShardingConfig` constructors (TP/SP/EP placements), used by `flavors.py`; omitted when the shared ones in `models/common/` suffice - `state_dict_adapter.py` -- HF checkpoint mapping - `__init__.py` -- re-exports the model class, `MODEL_FLAVORS`, `build_model_config` Training recipes live in `torchtitan_recipes/`, not in the model folder. ## Base Classes vs Shared Components Abstract contracts live in `torchtitan/models/base/`: `BaseModel` (model lifecycle), `Module` and its containers, `ShardingConfig`, and the state dict adapter bases. Concrete components used by multiple models belong in `torchtitan/models/common/`, including base architectures such as `Decoder`: - Attention mechanisms (`attention/`) - Feed-forward layers (`feed_forward.py`) - Decoder blocks (`decoder.py`, `decoder_sharding.py`) - MoE routing and expert implementations (`moe.py`, `moe_sharding.py`, `token_dispatcher.py`, `dist_moe/`) - Rotary/positional embeddings (`rope.py`) - Normalization layers (`norm.py`) - HF state dict adapter for MoE models (`moe_state_dict_adapter.py`) - Parameter counting and FLOP formulas (`flops.py`) ## Don't Over-Specialize If a feature is only needed by one model (e.g. HFTransformerModel), implement it in that model's folder. Don't modify shared infrastructure or base classes to accommodate a single model's needs. ## Control Flow in Forward Keep control flow (routing decisions, conditional logic) in the `forward` method. Don't bury important branching logic inside helper methods where it's hard to trace. --- name: cuda_graph_trace_compaction description: Make a CUDA-graph PyTorch profiler trace readable in Perfetto by packing replay streams into semantic lanes, and optionally merge per-rank pipeline-parallel traces into one aligned timeline with send/recv flow arrows. Use when a trace captured with CUDA graphs shows hundreds of one-kernel stream rows, or when the user wants to view several PP rank traces together. --- # CUDA Graph Trace Compaction CUDA graph replay gives each captured region its own stream, so a profiler trace opens in Perfetto as hundreds of near-empty rows. `scripts/compact_cuda_graph_trace.py` repacks the slices into a few semantic lanes (compute, NCCL all-gather, reduce-scatter, ..., memcpy, memset). **Visualization only.** Slices keep their name, `pid`, `ts`, `dur`, and `args` -- only the display placement (`tid`/`args.stream`) changes. Hand the user the original for measurement and the compacted copy for reading. The script is agent-maintained, not a torchtitan API -- edit it for the trace in front of you rather than working around it. Common tweaks: `_classify_slice` and `GROUP_ORDER`/`GROUP_LABELS` for new lane groupings, `PP_ANNOTATION_PATTERN` for a different annotation scheme. Keep the tests passing. ```bash # Single trace -> _compacted.json.gz beside the input python3 .claude/skills/cuda_graph_trace_compaction/scripts/compact_cuda_graph_trace.py \ outputs//profile_trace/iteration_10/rank0_trace.json.gz # All rank*_trace.json[.gz] in a directory -> pp_traces_merged_compacted.json.gz python3 .claude/skills/cuda_graph_trace_compaction/scripts/compact_cuda_graph_trace.py \ outputs//profile_trace/iteration_10 --merge-pp-ranks ``` `-o` sets the output path, `--overwrite` replaces an existing one. Single-rank compaction includes rank-local PP arrows (`UNSHARD -> compute`, `RECV -> compute`, and `compute -> SEND`). Cross-rank `SEND -> RECV` arrows are added only when merging PP ranks. Merging aligns ranks on a common `baseTimeNanoseconds`, stacks them in pipeline order, and draws flow arrows between the `PP:` annotations that `torch.distributed.pipelining` emits (`SEND_F -> RECV_F` across ranks, `RECV_F -> F -> SEND_F` within one). It also connects each `PP:UNSHARD` annotation on the first two all-gather lanes to the first subsequent same-stage compute annotation on the first compute lane. No annotations means no arrows, but the merge still works. The script prints a JSON summary and validates itself, raising if a measured slice changed or a lane ended up with overlapping slices. Worth a look: `original_streams` vs `compacted_lanes` for the compaction ratio, and `pp_send_recv_unmatched` (nonzero usually means a truncated capture window, not a pipeline bug). Rank order comes from a `torchtitan_real_pp` group or a `mesh_pp` group description in `distributedInfo.pg_config`; absent that it sorts by global rank, which is only correct when PP rank order matches global rank order. ## Tests `pytest .claude/skills/cuda_graph_trace_compaction/scripts/` -- hand-built traces, no GPU. Outside `testpaths`, so `pytest tests/` skips it. Run it after touching the script; a broken invariant yields a plausible-looking wrong trace rather than an error. --- name: fake_distributed_backend description: Run TorchTitan with fully fake communication or real pipeline communication plus fake SPMD axes, especially for logical-rank memory debugging. Use for fake process groups, real-PP/fake-SPMD runs, large-model dry runs, or comparing rank-local CUDA memory ownership without launching the full logical world. --- # Fake Distributed Backend TorchTitan provides two distinct test modes: - `fake` makes every distributed axis fake. Use it for configuration, shape, ownership, and rank-local allocation analysis without real transport. - `real_pp_fake_spmd` uses one physical process per pipeline rank. PP traffic is real; data, tensor, context, and expert parallel axes are fake. Use it when pipeline communication and buffer lifetimes must be exercised. Neither mode replaces a real distributed numerical or performance test. The authoritative user-facing launch contract, including complete command examples and environment-variable semantics, is in [`docs/debugging.md`](../../../docs/debugging.md#fake-backend-debugging). ## Logical Rank Coordinates Set `NGPU` to the logical world size. For pure fake PP, also set: ```bash export FAKE_PP_RANK= ``` For PP degree `P`, the non-PP logical world size is `NGPU / P` and the global logical rank is: ```text FAKE_PP_RANK * (NGPU / P) ``` Both modes represent SPMD coordinate zero. Pure fake defaults to PP coordinate zero when PP is disabled. In hybrid mode, physical `RANK` is the PP coordinate and `WORLD_SIZE` must equal the PP degree; setting `FAKE_PP_RANK` is invalid. `COMM_BACKEND=fake` is a convenience understood by `run_train.sh` for a single-process pure-fake run. Launch hybrid mode with `torchrun`, set `NGPU` to the logical world size, and pass `--comm-backend real_pp_fake_spmd`; `torchrun` supplies the physical `RANK`, `WORLD_SIZE`, `LOCAL_RANK`, and rendezvous variables. ## Memory Debugging Workflow 1. Use the production model, recipe, dtype, parallel degrees, microbatch count, FSDP policy, activation checkpointing policy, and CUDA-graph setting. 2. Choose the PP coordinate whose SPMD-zero ownership is being investigated. 3. Capture allocator state after initialization, after complete optimizer warmup, during steady-state forward/backward, and after optimizer completion. 4. Compare the same logical rank and capture point between configurations. 5. Escalate in order: pure fake -> real PP/fake SPMD -> real distributed run. Fake execution represents PyTorch-managed model parameters, optimizer state, prepared quantized weights, activations, gradients, pipeline buffers, and DistMoE scratch or activation arenas. It does not faithfully represent NCCL communicator allocations, network registration, transport scratch, collective latency, or communication overlap. Use `torch.cuda.memory_snapshot()` for allocation provenance and `torch.cuda.max_memory_allocated()` / `torch.cuda.max_memory_reserved()` for comparable totals. A fake run can establish deterministic ownership and lifetime behavior; it cannot establish distributed numerical parity or throughput. ## Failure Interpretation - A divisibility or coordinate error is a launch-contract failure; fix the logical topology rather than changing model configuration. - A hybrid world-size error means the physical allocation does not contain exactly one process per PP rank. - A pure-fake success followed by a hybrid failure usually isolates PP transport, send/receive ownership, or communicator initialization. - A hybrid success followed by a real-run failure usually isolates an SPMD collective, real communication memory, or scale-dependent scheduling issue. - Do not compare different PP coordinates as evidence of a memory regression; their model and activation ownership can differ intentionally. --- name: kernel-numerics-verifier description: Verify forward and backward numerical correctness of a handwritten low-precision GPU kernel against a high-precision oracle and an eager same-dtype baseline. Use for Triton or CUDA kernel correctness; use numerics_debugging to locate an unknown whole-model divergence. --- # Kernel Numerics Verifier Verify a handwritten kernel with three implementations: - `golden`: eager FP64 implementation. Use FP32 only when FP64 is unsupported. - `reference`: eager implementation in the kernel's input and output dtype. - `target`: handwritten Triton or CUDA kernel in the same dtype as `reference`. Generate low-precision inputs once, then cast those exact values to FP64 for `golden`. This measures kernel arithmetic error without mixing in input quantization error. ## Correctness gate For each output tensor and gradient tensor, compute in the golden dtype: ```text reference_error = max(abs(reference - golden)) target_error = max(abs(target - golden)) rounding_floor = max(abs(golden.to(target_dtype).to(golden_dtype) - golden)) absolute_floor = max(project_atol, rounding_floor) threshold = fudge_factor * reference_error + absolute_floor ``` When the reference is sufficiently accurate, pass when: ```text target_error <= threshold ``` Use an existing operator tolerance when one exists. Otherwise start with `fudge_factor = 2.0` and `project_atol = 0`. Declare them before running the target. Do not increase them merely to make a failing kernel pass. ## Check conditioning before gating Measure whether the reference error is material relative to the golden result: ```text golden_scale = max(abs(golden)) reference_relative_error = reference_error / max(golden_scale, absolute_floor) ``` Predeclare an acceptable reference-relative error. If none exists, use 10% as a diagnostic trigger. When the reference exceeds it, do not issue a pass or fail from the 2x gate; the permitted error is too large to distinguish a bug from numerical sensitivity. For a reduction, inspect the output element responsible for the error and capture its FP64 per-term contributions: ```text condition_number = sum(abs(per_term)) / abs(sum(per_term)) ``` A large condition number means cancellation amplifies small errors in the terms. It does not by itself prove that the target is correct. When the gate is invalid: 1. Compare two legitimate FP32 reduction orders. This measures sensitivity to accumulation order. 2. Repeat across several reduction lengths and seeds. Error growing roughly with `sqrt(T)` indicates zero-mean numerical noise; growth with `T` indicates systematic bias. Track the condition number at each `T`. 3. Compare reference and target error distributions and bias. A target-specific bias, abnormal scaling, or semantic mismatch is a failure. If both implementations follow the same conditioning-driven noise envelope, report `PASS WITH LIMITATIONS`. If the evidence cannot distinguish numerical noise from a kernel bug, report `INCONCLUSIVE`. For simple elementwise kernels, also require the result to be correctly rounded or within 1 to 2 ULPs. Larger factors require evidence from an already accepted implementation across representative inputs. ## Workflow 1. Read the kernel, eager formula, callsites, and existing tests. Record the formula, intermediate dtypes, accumulator dtypes, output dtype, supported shapes, masking, and NaN/Inf behavior. 2. Create identical low-precision inputs for `reference` and `target`. Create `golden` inputs by losslessly upcasting those values. 3. Compare shape, dtype, NaN locations, and positive and negative Inf locations before applying numerical tolerances. Any mismatch is a failure. 4. Check the reference-relative error. Apply the correctness gate only when the reference is sufficiently accurate; otherwise run the conditioning analysis. 5. Apply the resulting decision procedure to every forward output. 6. Use one random upstream gradient for all three implementations. Check every input and parameter gradient separately. 7. Test production shapes, the smallest legal shape, and dimensions immediately below and above relevant kernel tile sizes. Add adversarial values relevant to the operation, such as cancellation, extreme logits, or near-zero variance. 8. If the kernel is nondeterministic, repeat identical inputs and require every result to pass the same gate. 9. Diagnose failures before changing tolerances. Check indexing and masks first, then accumulator dtype, premature casts, reduction order, numerical stabilization, approximation functions, and backward formulas. ## Result Report a compact table for forward and each gradient: | Tensor | Ref/golden | Condition | Target error | Threshold | Gate valid | Verdict | | --- | ---: | ---: | ---: | ---: | --- | --- | Give one verdict: `PASS`, `PASS WITH LIMITATIONS`, `FAIL`, or `INCONCLUSIVE`. State the tested shapes, dtypes, and unsupported cases. Do not claim general correctness from one shape or forward-only testing. Benchmark performance only after numerical verification passes. # What torch.compile can't do today Limits with a known rewrite are in symptoms.md and patterns.md. These are the rest: design around them. Format: **limit**: symptom -> workaround. - **FMA contraction**: compiled differs from eager by ~1 bf16 ulp (seen on GB300; H100 was bitwise) -> tolerant tests (~1 ulp); bitwise only for batch invariance. - **Masked loads with symbolic T (GB300)**: many-load kernels ~2x slower -> keep small dims static; T stays dynamic (one graph), so expect this on broadcast-heavy kernels. - **`maybe_mark_dynamic` forbidden in graphs**: trace error -> mark in eager before the region (pattern 5). - **size-0/1 dims always specialize**: extra graphs -> avoid size-1 call signatures, or budget for them. - **make_fx specializes Python ints and captures class-attribute tensors as constants**: stale values under GraphTrainer -> graph inputs, or update a persistent tensor in place. - **Region options dropped under an outer compile**: an option has no effect under block compile or GraphTrainer -> express the fix in code, not only as an option. - **Functional regions can't alias outputs or update grads in place**: extra copies -> hand-written override when it matters. - **SAC replays regions; saving them (`wrap_inductor_compiled_regions`) hits a storage bug after a recompile**: SAC recompute cost, or a crash -> measure under SAC; region-level SAC (gating-and-options.md). - **Store to a new attribute of an outer-scope object inside a checkpoint HOP**: generic "Observed exception" -> create the state outside (e.g. a thread-local's `__init__`). - **`autograd.Function` reading `w._base` / `grad_dtype` of a weight view**: AOTAutograd crash -> avoid in traced code. - **Mesh push/pop inside traced regions**: hard error -> keep mesh context outside regions. - **Tensor caches (`functools.cache`, module-level dicts) under FakeTensorMode**: fake tensors cached and reused in real runs -> key caches on real tensors only, or skip under fake mode. - **Inductor's FX graph cache ignores a changed custom-op fake**: the old (wrong) fake is still used -> clear the cache dir after changing a fake. A new limit hit during a run: add it here (symptom, workaround, a minimal repro command), and report it per pr-template.md. Don't file it outside the repo without asking the user. # Results that change when the setup changes Each row fooled us or nearly did. Before claiming a win, check the rows that apply. | Dimension | What flipped | Check | |---|---|---| | Fixes combined | DSv3 `mla_qk`: real-arithmetic RoPE alone slower (46.95 ms), cat option alone neutral (46.37), both 46.34 -> 45.85 ms | before rejecting a fix, try it with the others | | TP=1 harness vs a TP>1 recipe | a one-GPU harness at TP=1 is only a proxy when the recipe runs TP>1 (GPT-OSS RL: TP=2, EP=2) | say so in the numbers; EP-only ops: harness.md (EP work on one GPU) | | TP=1 vs TP>1 | folding the residual add into the MoE combine looked free at TP=1 but is wrong under TP (the output is partial before its all-reduce) | the sharding of every tensor a rewrite touches | | TP=1 vs TP>1 | a per-row score scale before w2 is exact only because routed w2 isn't sharded across the row | same | | World 1 vs EP>1 | zero-token ranks exist only at EP>1: HybridEP's combine backward skipped probs there, so router grads went missing silently | a real multi-GPU run for cross-rank paths | | World 1 vs EP>1 | dense top-k input and a fused permute were neutral or slower at world 1; their possible wins appear only at scale | measure at the target scale before rejecting for good | | H100 vs GB300 | masked-load reductions ~2x slower on GB300; FMA changed bits on GB300 (an upstream CosSinRoPE test failed only there); bandwidth ~7.1 vs ~3.35 TB/s; attention backend rankings (cuDNN vs FA4) differ | label hardware; don't transfer rankings | | Config default vs recipe flags | blocking dispatch (the HybridEP default) had a wrong fake and the zero-row bug while every test recipe used non-blocking; a win reached only test recipes because no production recipe ran HybridEP | which configs actually reach the change | | Activation checkpointing | no-AC wins shrink or flip under SAC (regions replayed, ~20 µs CPU per eager op) | measure with the recipe's AC mode (gating-and-options.md) | | Eager vs compiled vs GraphTrainer | a rewrite's sign can differ per mode and per model | negative-results.md; measure each mode | | Missing optional deps | DSv4's venv lacked FA4 and had attn-gym 0.0.13 (0.0.16 required): `gather_attn` silently fell back to Triton and attention took ~115 of 180 ms | check attention kernel names in the profile; pin optional deps. Also: varlen attention on Blackwell needs FA4 | | Attention backend | a head-major q / o-grad layout was -3.8% with Triton attention and neutral with FA4; deterministic mode switches backends | measure layout changes with the production backend and mode | | Small shapes | host-bound: kernel time and wall time disagree | report both | | Harness | detached inputs, accumulated grads, unseeded buffers, kernels merged by name, order effects, shared GPU, timing a harness copy | measurement.md (traps) | | Cold compile and graph budget | a ~2% layer win cost ~3 min of cold compile per start; the limit was often blown only by the first `no_grad` pass or an untested shape | report the trade-off; count all shapes | # Gating, options, numerics and parallelism ## Which code path runs where A gate is an if-branch that picks the rewritten form only when compiling. A rewrite that helps compile can hurt eager; the eager cost (measured with inputs requiring grad) was our most useful guide. Defaults, not rules: ```text new form slower in eager? no -> one path, no gate (simplest; also helps GraphTrainer) yes -> does it also win under GraphTrainer full Inductor? yes -> compile_friendly() (if this tree has it; else is_compiling()) no / model-dependent -> torch.compiler.is_compiling() Dynamo-only hints (mark_static, maybe_mark_dynamic) -> is_compiling() ``` Example of `compile_friendly()`: MoE dispatch/combine as a gather-sum costs +3.8 to +4.1 ms per 2 layers in eager (screen) but wins under GraphTrainer. The verdict can flip per model and Inductor config: the ComplexRoPE rewrite was slower under GraphTrainer on two models with default cat lowering, and faster on one of them with ConcatKernel. ## Region options `@local_compile(name, batch_invariant=..., options={...})` passes options to that region's own `torch.compile`. They're dropped when the region is inlined under an outer compile (block compile, GraphTrainer). | Option | What it does | Caveat | |---|---|---| | `max_pointwise_cat_inputs=0`, `max_complex_pointwise_cat_inputs=0` | per-input cat copies (ConcatKernel) | wins on some shapes, loses on others (patterns.md 2) | | `unroll_reductions_threshold=K+1` | unrolls a static small reduction | no effect on symbolic sizes | | `triton.mix_order_reduction_non_strict_mode=True` | no guards from the mix-order heuristic | time-neutral in our cases | | `coordinate_descent_tuning=True` | autotunes reductions (HcPost backward 739 -> 594 µs) | doubles cold compile, can slow decode shapes; turn off when `torch.are_deterministic_algorithms_enabled()` (Inductor only skips it on its own `config.deterministic`) | | `triton.prefer_nd_tiling=True` | tiles `[T, N, D]` as rows x D so a symbolic small N isn't divided into every index | Kimi dynamic-width backward 1409 -> 571 µs | | `assume_32bit_indexing=True` | int32 index math | hard compile failure once T*N*D > int32; we rejected it | ## Numerics - Compiled fused pointwise math is not bitwise vs eager: Triton contracts into FMAs, so ~1 bf16 ulp differences appear. Tests should allow ~1 ulp; keep bitwise checks only for batch invariance. - Routing-dependent tensors (router gate grads) can differ by 0.1 relative under any compile because ties in top-k flip with rounding. Check that eager-vs-eager variants and block compile show the same pattern before blaming your change. - Report rel L2 vs fp64; fused fp32 accumulation often makes compiled closer to fp64 than eager. - Correctness procedure: measurement.md. ## Activation checkpointing It matters, so measure with the AC mode the recipe uses. - **Activation remat (current main, `torch_remat`):** code outside declared remat regions is always replayed in the backward, including eager glue around `local_compile` regions. A bare consumer of a remat region's output needs `remat.recompute_needs_tensor`. Check how your compile regions line up with the remat regions: a compile region that falls outside a remat region is replayed. Nest compile inside remat, never remat inside a compiled region. Older trees (eager SelectiveAC): - **Full AC:** each block's forward runs twice. Region forwards run twice too, so a region's forward win counts double, and so does any eager glue left outside regions. - **Selective AC (eager):** saves a policy-chosen set of op outputs and recomputes the rest. A `local_compile` region's outputs are not on the save list by default (`torch._inductor.config.wrap_inductor_compiled_regions` is False), so regions are replayed in the backward. Turning that flag on crashed MoE models after a recompile, because SAC keys saved outputs per compiled cache entry. - **SAC host overhead:** SAC's dispatch mode costs ~20 µs of CPU per eager op. With many small eager ops (e.g. 4k tokens per rank), the step becomes host-bound: on DSv3 at 4k, regions under SAC were 48.6 ms vs 34.3 ms for block compile + SAC, mostly host time. - **Region-level SAC:** calling `torch.utils.checkpoint` with the SAC policy inside a compiled half-block applies the policy at trace time with no runtime dispatch mode; it matched block compile + SAC in our prototype (needs care with side effects and thread-local state inside the checkpoint). ## Parallelism checklist - **FSDP/HSDP:** regions see local shards; no change unless a region reads parameters that FSDP hooks touch. - **TP:** is the rewrite exact per shard? A per-row scale before w2 is exact when w2 is replicated over TP (routed experts under EP); a residual add folded into a TP-partial output is wrong before the all-reduce. - **EP:** shapes are rank-local; zero-token ranks happen at large EP. Check fakes and backward for empty inputs, and that collectives take the same code path on every rank. - **PP:** stage metadata is static; regions are fine. - **CUDA graphs:** blocking ops with host syncs or data-dependent sizes cannot be captured; check capture and replay bitwise. - **GraphTrainer / make_fx:** class-attribute tensors become constants (update a persistent tensor in place); Python ints specialize; wrong fakes record wrong shapes. - **Thread-local state:** storing any new attribute on an object defined outside a checkpointed region (e.g. lazily creating a list on a module-level `threading.local`) fails under Dynamo with a generic "Observed exception"; create the state outside (e.g. in the thread-local's `__init__`). # Harness and checks (copy-paste) ## Where regions live - Declare: `@local_compile("name", batch_invariant=..., options={...})` on a function (`torchtitan/distributed/local_compile.py`). `grep -rn "@local_compile(" torchtitan/` lists existing names. - Enable: the model config's `local_compile_regions` list (e.g. `["loss", "swiglu"]`), set in the model's config builder or appended in a recipe (`config.local_compile_regions.append("router")`). - Bind: `apply_local_compile(regions)` (the trainer calls it); in a harness call it yourself after building the model. - The API differs by torchtitan version. Newer main keeps regions in the model config (`local_compile_regions`). Older trees set them at recipe level, `Trainer.Config.compile = LocalCompileConfig(regions=...)` with a shared default list, bound with `.apply_local_compile()`. Check which applies: `grep -rn "LocalCompileConfig\|local_compile_regions" torchtitan/`. ## Bandwidth (once per GPU) ```python x = torch.empty(2**30, dtype=torch.uint8, device="cuda"); y = torch.empty_like(x) for _ in range(3): y.copy_(x) s, e = torch.cuda.Event(True), torch.cuda.Event(True) s.record(); [y.copy_(x) for _ in range(20)]; e.record(); torch.cuda.synchronize() tb_s = 2 * x.numel() * 20 / (s.elapsed_time(e) / 1e3) / 1e12 # read + write ``` ## 2-layer step at native dims ```python torch.manual_seed(0) torch.backends.cuda.matmul.fp32_precision = ... # whatever this tree's trainer sets (grep fp32_precision torchtitan/) cfg = ModelConfig(...) # the model's real flavor config blocks = nn.ModuleList([Block(cfg, layer_id=i) for i in (a, b)]).cuda().to(torch.bfloat16) # pick layer ids that cover each block type (e.g. one dense + one MoE layer) apply_local_compile(regions) # [] for eager, the upstream list, or your variant x = torch.randn(tokens, cfg.dim, device="cuda", dtype=torch.bfloat16, requires_grad=True) # read Model.forward for what each block takes (positions, rope cache, masks); build those once here, # with requires_grad where training has it (e.g. a residual stream passed between blocks) extra = dict(...) def model(x, **extra): for block in blocks: x = block(x, **extra) return x def step(): for p in blocks.parameters(): p.grad = None x.grad = None out = model(x, **extra) out.float().square().mean().backward() for _ in range(5): step() # warmup (compiles) torch.cuda.synchronize() s, e = torch.cuda.Event(True), torch.cuda.Event(True) s.record() for _ in range(20): step() e.record(); torch.cuda.synchronize() wall_ms = s.elapsed_time(e) / 20 ``` - When production calls a region with a dynamic small dim (e.g. a stack width that varies per layer), warm the harness at several widths first; otherwise it times a static graph production never runs. - MoE layers with an aux loss need `AuxLoss.set_step_denominator(...)` before the first forward, as the trainer does. - Run each variant in a **fresh process** for cited numbers; one process with interleaved variants for screening. - Record `torch.cuda.max_memory_allocated()` (reset before the timed loop) and `nvidia-smi --query-compute-apps=pid` before and after. - **EP work on one GPU:** a harness with expert parallelism off silently skips the dispatcher's EP-only ops (e.g. `_permute` / `_unpermute`; on DSv3 the pair cost 9.36 ms per MoE layer at 16k at EP=2 shapes, 4.08 after a rewrite). Benchmark EP-path ops on one GPU at a real rank's shapes, with a mocked EP mesh or by calling the ops directly. - **MoE on one GPU:** - local (no-EP) dispatcher, with local experts set to one EP rank's share (e.g. 256 experts / EP8 = 32) so grouped GEMM shapes match a real rank; - the router then scores fewer experts, so also time the router at full width; - libraries that need a process group: a 1-rank group for timing; a fake process group (`torch.testing._internal.distributed.fake_pg`) only for FakeTensor tracing (fake peers alias storage). - Benchmark the code you will ship, not a harness copy of it: the shipped version may compute extra things (an in-graph inverse permutation cost ~18 µs the copy didn't). - Before a multi-GPU cluster job, build the model at the target world size on meta tensors or a fake PG first; a world-1 smoke run won't catch config checks that fire only at EP > 1. ## Profiler kernel time ```python from torch.profiler import profile, ProfilerActivity with profile(activities=[ProfilerActivity.CUDA]) as prof: for _ in range(5): step() torch.cuda.synchronize() rows = [(e.key, e.device_time_total / 5, e.count / 5) for e in prof.key_averages() if e.device_time_total > 0] rows.sort(key=lambda r: -r[1]) kernel_ms = sum(r[1] for r in rows) / 1e3 for name, us, n in rows[:40]: print(f"{us:9.1f} us x{n:4.1f} {name[:110]}") ``` `key_averages` groups by name: identical kernels merge, so look at the count column. ## Mapping Triton kernels to source - Inductor names kernels after the fused ops: `triton_poi_fused_add_mul_silu_1` (pointwise), `triton_red_*` (looped reduction), `triton_per_*` (persistent reduction). - `TORCH_LOGS=output_code` prints the generated code with the source line of each fused node; `TORCH_COMPILE_DEBUG=1` writes it to disk. - `TORCH_LOGS=graph_breaks,recompiles` shows why a region breaks or recompiles (guard failures name the shape or global state). ## Counting graphs ```python torch._dynamo.reset() from torch._dynamo.utils import counters # run every shape the model uses: each training token count, each call site, then a no_grad pass print(counters["stats"]["unique_graphs"]) ``` Or for one function: `torch._dynamo.eval_frame._debug_get_cache_entry_list(fn.__code__)` (one entry per graph). The limit is `torch._dynamo.config.recompile_limit` (8); with `fullgraph=True`, hitting it raises. ## Eager cost and GraphTrainer timing - Eager cost of a rewrite: force the new branch on in eager (patch the gate) and time the 2-layer step with inputs requiring grad. - GraphTrainer (only if the model trains with it): trace the 2-layer step with make_fx and compile it with `full_inductor_compilation_pass` (`torchtitan/experiments/graph_trainer/inductor_passes.py`), then time old vs new branch in the same process from profiler kernel time. ## CUDA-graph and SAC checks ```python g = torch.cuda.CUDAGraph() s = torch.cuda.Stream(); s.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(s): for _ in range(3): step() # warm up on a side stream torch.cuda.current_stream().wait_stream(s) with torch.cuda.graph(g): step() g.replay(); torch.cuda.synchronize() # compare outputs/grads to an uncaptured step ``` Steps that zero grads inside the capture run extra kernels; compare captured vs captured. For SAC, apply the model's AC the way the trainer does and time again. On main (`torch_remat`), checkpointing blocks without configuring their region names fails ("Duplicate torch_remat region name"); call `config.activation_remat.apply` on a holder module like the trainer. On older trees eager SAC replays region forwards in the backward. ## Environment gotchas - `python -c` run from the repo root imports the main checkout, not your worktree on `PYTHONPATH`: run harnesses from another directory. - `git stash` and branch names are shared across worktrees; use WIP commits instead of stashes. - Subagents' file-writing tools may refuse report files; writing them via the shell works. # Measurement Most of the mistakes in this work were measurement mistakes, not code mistakes. This file is the procedure plus every trap we fell into. ## Ideal time For a memory-bound op: ideal = (bytes of every input read once + every output written once) / measured bandwidth. - Measure bandwidth once with a large copy (GB300: ~7.1 TB/s; H100: ~3.35 TB/s). Spec-sheet numbers overstate what kernels reach. - Count fwd and bwd separately. A pointwise op's backward reads the grad and saved inputs and writes input grads. - Example: multiplying `[32768, 7168]` bf16 by a per-row score reads 470 MB and writes 470 MB: 940 MB / 7.1 TB/s = 132 µs. The eager broadcast kernel took 507 µs (3.8x), which made it the first target. - For GEMMs and attention use a FLOP floor and the achieved rate of the best kernel you know (GEMMs reach ~1.8-2.0 PFLOP/s bf16 on GB300). - The best known kernel for the same op is a second, reachable target: report the gap to both. Ideal is a floor nobody reaches; a hand-written kernel proves what is reachable and its profile shows which kernels, copies or saved tensors differ. - Attribute a gap fully: per-kernel times for both sides, kernel sum + launch gaps = wall. "Not explained yet" means the diff wasn't done. Chasing parity still follows the ~1% bar: list what's left instead of adding code for 0.3%. - "Gap to ideal × share of the layer" ranks candidates. A 3x gap on 1% of the layer matters less than 1.3x on 15%. ## Correctness - Rel L2 of the output and every grad vs eager, run once in bf16 and once in fp32, plus an fp64 reference. Do the fp32 and fp64 checks at op level: an fp32 layer-level run with flex autotuning took over 30 min. - Tests that fail before a perf change: region-vs-eager parity at op level (within ~1 bf16 ulp); traced placeholder shapes are concrete when a dim is meant to be static (`mark_static`); graph-count assertions over the shapes the model uses. - MoE layers: routing flips make raw comparisons noisy. Replay the eager routing in the compiled run (patch the model's own router class: DSv4 overrides `_select_experts`, so patching the base class did nothing) and compare against a noise band from a 1-ulp perturbation of the input. - Don't dump full MoE grads for comparison (53 GB per dump on DSv4); compare per-tensor norms, samples or hashes. - Op-level first (isolates the kernel), then layer-level (neighbors, host time, memory). ## Traps | Trap | What happened | Fix | |---|---|---| | Detached step inputs | the residual stack needed no grad, so one form looked 2 ms faster in eager when it was 5.75 ms slower | inputs require grad as in training; treat any `.detach()` on step inputs as suspect | | Grads accumulate across steps | bf16 accumulation adds, ~1.3 ms per DSv3 MoE layer, inflated times and "% of layer" | set grads to None every step | | Unseeded `torch.empty` buffer | cross-process accuracy comparisons were meaningless | seed before building inputs and weights | | Summing kernels by name | two 7 µs scatters showed as one 14 µs entry | check kernel counts before attributing | | "fp64 reference" calls `.float()` inside | silently computed in fp32 | check the reference's dtypes | | Routing replay keyed by call count | broke under AC recompute (the router runs again) | key by layer and step | | fp32 matmul precision differs from the trainer's | fp32 router/gate GEMMs ran as SIMT kernels; wrong baseline | set `torch.backends.cuda.matmul.fp32_precision` to what this tree's trainer sets | | Another job hits mid-run | a run that started clean showed a 2x outlier wall; start/end `nvidia-smi` checks miss short jobs (one landed inside a timed loop) | poll `nvidia-smi` during the run; report the median of processes within ~1% of the best, with n shown | | A crashed run reusing the last result | the harness printed the previous RESULT line after a crash | write a FAILED marker per run | | Order effects | a variant always run second carried warm caches or a slower host | rotate order or use fresh processes | | Pooled batches | pooling a harness batch with a PR-tree batch hid a 1.4% vs 0.3% disagreement | report batches separately | | Timing a harness copy | the shipped code computed an extra in-graph inverse permutation (~18 µs) | time the code you ship | | World-1 smoke before a cluster job | a config check that fires only at EP > 1 killed every rank | build at the target world size first | ## Profiling and trace diffs - Profile eager, regions and block compile on the same inputs. Group kernels by family (GEMM, attention, norm, cat/copy, pointwise, reduction, scatter/gather, sort). - The regions-vs-block diff is your list of glue that a region could still take. Typical leftovers: cat kernels (`CatArrayBatchedCopy`), broadcast muls, `index_put`/scatter backward, casts, standalone adds. - Before calling a kernel "outside the region", check its parent in the trace (the `Torch-Compiled Region` range). An ATen kernel inside a region is an Inductor fallback (complex mul, `topk`, custom ops): it needs a rewrite (symptoms.md), not a bigger region. We once missed RoPE's complex mul (2 x 212 µs) by calling it outside. - When two fixes each look neutral, try them together: rewrites and Inductor options interact (patterns.md 1). - Look for host syncs: `torch.cuda.set_sync_debug_mode("warn")`, `cudaStreamSynchronize` and DtoH memcpy in the trace. Boolean-mask indexing, `.item()`, `.tolist()`, `repeat_interleave` without `output_size`, and blocking copies from pageable or pinned memory all sync. - Always measure CUDA-graph replay too (whole fwd+bwd captured after warmup): replay ≈ kernel time + a few µs per kernel. If a variant can't be captured (host sync, data-dependent shape), say why. ## Replicates and hygiene - **On a shared GPU,** wrap each fresh-process run in a wait-for-quiet-GPU + retry loop (poll `nvidia-smi`, rerun contaminated rows); that gave clean numbers. When the GPU never goes quiet, interleaved A/B inside one process still gave correct relative deltas under 2.3x contention. - **Screen in one process, confirm in three.** In-process A/B (variants interleaved) is good for ranking. Cited numbers are medians of 3 fresh processes. - **Busy shared GPU and many branches:** a full fresh-process matrix per branch may not fit. A fallback we used: clean 3-process medians for main vs the final stack, plus single-run per-branch deltas labelled "indicative". - **Host-bound regimes.** At 4k tokens per rank, eager plus many small kernels can be host-bound; walls swung 1.5-3x with other agents' CPU load. Report GPU kernel time alongside wall, and an interleaved in-process A/B to show a region's calls pay for their host cost. - **Record per row:** wall and kernel time, kernel count, peak memory, graph count, foreign GPU processes, commit hash. - **Label everything:** shape, fwd+bwd or fwd, wall or kernel, process count, GPU. "13% faster" without these is not a claim. ## Graph budget - Count graphs for every shape the model calls the region with (all layers, routed vs shared experts, different head counts), for every training token count, and for a `no_grad` pass (grad mode is a guard, so validation adds a graph per shape class). - `recompile_limit=8` with `fullgraph=True` is a hard failure, not a fallback. Headroom helps (we aimed for ≤ 6 of 8). Raising the limit (we used 16 in several models) costs more compiles and more cached graphs, and how to do it depends on the tree: - newer upstream main: a per-region `recompile_limit` where the region supports it (e.g. `fused_binary_activation`); - trees that pass `local_compile` kwargs straight to `torch.compile`: there is no such argument there; - otherwise: `torch._dynamo.config.recompile_limit`, globally. - Graph count also tells you whether a symbolic size crept in (a static graph followed by a dynamic one). - Every new call signature costs one graph per grad mode. - A region shared by call sites of different rank or layout multiplies graphs; give call sites stable layouts or separate regions. - `triton.mix_order_reduction_non_strict_mode=True` stops a scheduler heuristic from installing guards on symbolic sizes (worst case 8 -> 6 graphs, time-neutral). - Keep fixed dims static (`mark_static` inside the region) and mark the token dim dynamic in eager before the region (patterns.md 5). - How to count: harness.md. ## Cold compile - Measure with caches disabled (`TORCHINDUCTOR_FORCE_DISABLE_CACHES=1`), per graph. Autotuning options (coordinate descent) can double compile time; an unrolled loop can add minutes (a per-entry Sinkhorn rewrite added ~3 min cold). - Report it next to the win, as a trade-off ("~11 s more cold compile for 0.5% of step"; DSv4 per-entry unrolling: ~+42 s for 2% of step, stack total 48.8 -> 159.7 s) and let the owner decide. # Ideas that measured wrong Check these before spending time on them; each was measured, not guessed. Each was tested alone: it may still win combined with another fix (SKILL.md step 5). - **Per-site RMSNorm regions** on 7168-wide norms: ~0.4 ms per layer (0.7%), and at 4k the region's host time (330 vs 141 µs per norm) exceeded the kernel saving. - **Fused residual-add + RMSNorm region:** -0.7% at 16k, slightly slower at 4k (two more region calls per layer cost +0.35 ms host). Cross-block fusion needs a block restructure. - **Folding the residual add into the MoE combine:** wrong under TP (the MoE output is partial before its all-reduce). - **Global cat-lowering option:** faster for one RoPE, 24-26% slower for another; scope per region. - **Forcing unrolling via a big global threshold** (the first thing people try): has no effect on symbolic sizes. - **`assume_32bit_indexing` on large regions:** hard failure once T·N·D > int32. - **Coordinate descent on every region:** +2x cold compile and decode regressions; we kept it to the one region where it measured a clear gain, and turned it off in deterministic mode because its benchmark-picked configs can differ run to run. - **Gathers instead of scatters, assumed to win compiled too:** EP permute/unpermute as gathers was 2.7x in eager. Compiled, the permute half still gained and the pair was a net win, but the unpermute half was 7% slower (Inductor's scatter was already fast and the in-graph inverse permutation cost ~18 µs). Measure each half compiled. - **A library's fused permute+dispatch:** bitwise but slower than separate kernels. - **Passing dense top-k ids instead of a routing map:** bitwise, not faster within one node (the win is a smaller allgather at large scale). - **Removing a host sync:** worth it for overlap and CUDA graphs, but often within noise for wall time at one rank; frame it as a sync removal, not a speedup. - **Making a rewrite the single path because it won under compile:** check eager with inputs requiring grad; one "single path" verdict reversed (+5.75 ms eager) once the harness stopped detaching inputs. - **Wrapping compiled regions for SAC** (`wrap_inductor_compiled_regions`): crashed MoE models after recompiles (SAC storage keyed per cache entry). - **Block compile as a deliverable:** a measurement baseline only (not a ceiling: rewrites plus regions beat it); production configs use regions or CUDA graphs. # Rewrite patterns Each pattern: what blocks Inductor, the rewrite, and what we measured (GB300, fwd+bwd unless noted). Numbers are examples from DSv4, Kimi K3, DSv3 and Qwen3.5-MoE; the patterns are general. | # | Pattern | Typical win | |---|---|---| | 1 | complex ops -> real arithmetic | RoPE 443.7 -> 199.0 µs | | 2 | cat / stack lowering per region | 2155 -> 781 µs (or the reverse) | | 3 | tiny bmm -> explicit sum | 3475 -> 2155 µs | | 4 | small reductions: unroll | 1354 -> 290 µs | | 5 | keep small dims static; unbind in region | SwiGLU 1.68-1.75x | | 6 | gather-sum instead of scatter | combine 13540 -> 1102 µs | | 7 | scale the narrower tensor in a region | MoE module -12.8% | | 8 | int16 sort keys | 72.9 -> 30.4 µs | | 9 | per-entry unroll of tiny loops | 309 -> ~100 µs | | 10 | router glue as one region | 356.7 -> 135.9 µs | | 11 | data-dependent shapes, fakes, host syncs | correctness, syncs | | 12 | graph budget: moved to measurement.md | | | 13 | fp32 GEMMs, exact bf16 splits | 3262 -> 1291 µs; 0.288 -> 0.065 ms | | 14 | custom ops with exact fakes | SAC block 302.5 -> 122.6 ms | | 15 | loss path | -13% / -18% | | 16 | `autograd.Function` with fwd + bwd regions | 24.3 ms -> 1.56 ms per call | | 17 | small symbolic reduction -> GEMM | 1435 -> 180 µs | | 18 | input grad in the next kernel's layout | -3.8% (Triton attention) | | 19 | per-expert bias: gather + one-hot GEMM grad | GPT-OSS -9.7% | | 20 | eager hooks around attention as a region | GPT-OSS -3.9% | ## 1. Complex ops -> real arithmetic Inductor has no codegen for complex multiply. It keeps the ATen complex kernel and cannot fuse the split/cat glue around it ("Torchinductor does not support code generation for complex operators", printed only on a cold compile). ```python # before: complex RoPE x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) out = torch.view_as_real(x_c * freqs_cis).flatten(-2) # after (compile branch): same math on the real view of the cache cos, sin = torch.view_as_real(freqs_cis).unbind(-1) # [..., D/2] each x_even, x_odd = x.float().unflatten(-1, (-1, 2)).unbind(-1) out = torch.stack((x_even * cos - x_odd * sin, x_even * sin + x_odd * cos), dim=-1).flatten(-2) ``` DSv4 RoPE: compiled complex 443.7 µs (7 kernels) -> real arithmetic 199.0 µs (2 kernels). Compiled real arithmetic is not bitwise vs eager complex (FMA, see gating file). At strided call sites (q_pe sliced from a wider q) the real form's backward got slower than compiled complex (DSv3 671B: 1145 vs 1003 µs); a per-half cast plus the ConcatKernel option (pattern 2) brought it to 457 µs. Try the rewrite together with the cat option, not each alone: on DSv3's `mla_qk` region the real form alone was slower (46.95 ms) and the option alone neutral (46.37), but both together beat the region with complex left inside (46.34 -> 45.85 ms per layer at 16k). ## 2. Cat / stack lowering Inductor lowers a cat either as one masked pointwise kernel (per-element index math; fuses into consumers) or as per-input copies (`ConcatKernel`). Neither is right everywhere. - Inverse RoPE on `[16384, 64, 512]`: default 2155 µs, ConcatKernel 781 µs (ideal 605). - Norm + RoPE: default 1026 µs, ConcatKernel 1502 µs (the default lets the cat fuse into the norm's reduction). Neither lowering won everywhere for us, so it's worth trying both for the region in question; the option is region-scoped: ```python @local_compile("inverse_rope", batch_invariant=True, options={"max_pointwise_cat_inputs": 0, "max_complex_pointwise_cat_inputs": 0}) def inverse_rope(...): ... ``` A masked cat or stack consumed by a persistent reduction re-loads the row once per branch; avoid stacking inside a reduction, or split the reduction. ## 3. Tiny batched matmul -> explicit sum `[T, 1, 8] @ [T, 8, 7168]` goes to cuBLAS GEMV as an extern, so the bf16 -> fp32 upcast of the big operand is materialized (1.9 GB at 16k) instead of fused. ```python # before out = torch.matmul(probs.unsqueeze(1), values.float()).squeeze(1) # after: Inductor fuses the upcast into the reduction out = (probs.unsqueeze(-1) * values.float()).sum(1) ``` Kimi residual, T=16384: 3475 -> 2155 µs (ideal 794). Inductor's `decompose_mm_pass` does the same only above 10240 rows and not with BF16x9 fp32 matmuls enabled. ## 4. Small reductions: unroll instead of reduce Inductor unrolls a reduction into a pointwise kernel only when its size is a **static** int below `unroll_reductions_threshold` (default 8, strict `<`). Otherwise it emits a reduction over a strided dim that runs far below bandwidth. - Sum over K=16 rows: persistent reduction 1354 µs; Python loop over K (unrolled) 290 µs (ideal 281). - Weighted sum over N=8: persistent 1574 µs; threshold 9: 353 µs. ```python # after: a Python loop over a small static K becomes one pointwise kernel acc = rows[:, 0].float() for k in range(1, top_k): acc = acc + rows[:, k].float() ``` Or scope `options={"unroll_reductions_threshold": K + 1}` to the region. If the small size is symbolic, no threshold applies: see pattern 5 (keep it static), an `autograd.Function` that handles the width explicitly (Kimi residual dynamic width 3642 -> 2652 µs), or pattern 17. ## 5. Keep small dims static Automatic dynamic shapes make a dim symbolic after it takes a second value. A symbolic hidden size or stack width turns unrolled kernels into loops and adds int64 index math. ```python @local_compile("swiglu", batch_invariant=True) def swiglu(x_gate, x_up): torch._dynamo.mark_static(x_gate, -1) # hidden size F is fixed per call site ... ``` - **Check first:** newer trees already ship this for SwiGLU/SiTUGLU (`fused_binary_activation`: packed unbind, `mark_static`, recompile limit 16). - **`mark_static` inside the region:** SwiGLU with static F made dense/routed/shared calls 1.68-1.75x faster, at 5 graphs instead of 3. Marking a view in eager before the call did not stick. - **Unbind a packed projection inside the region:** pass `[R, 2, F]` and unbind in the region, so the unbind's backward (a `stack`, ~90 µs per MoE layer) fuses instead of running as a separate cat kernel. - **A Python loop over a dim forces specialization** even under automatic dynamic shapes, one way to keep a small width static. - **Views get symbolic strides and offsets** even when sizes are static; a packed `[rows, 2, F]` layout avoids it. - **Mark the token dim dynamic in eager**, before the region (`maybe_mark_dynamic` is forbidden inside a graph), guarded by `is_compiling()` so the marks don't run when the caller is itself compiled: ```python if not torch.compiler.is_compiling(): torch._dynamo.maybe_mark_dynamic(x, 0) # T varies torch._dynamo.mark_static(x, -1) # normalized dim never varies ``` ## 6. Deterministic gather-sum instead of scatter / atomics / opaque custom ops A custom op used to force determinism is opaque to Inductor and blocks fusion. When every output row has a fixed set of contributors (each token owns exactly K expert rows), write combine as a gather-sum: ```python # combine: out[t] = sum_k scores[t, k] * expert_out[inv_perm[t, k]], accumulated in fp32 rows = expert_out[inv_perm.flatten()].view(T, K, D) acc = rows[:, 0].float() * scores[:, 0, None] for k in range(1, K): # unrolled over K (pattern 4) acc = acc + rows[:, k].float() * scores[:, k, None] out = acc.to(expert_out.dtype) ``` Wrap it in an `autograd.Function` whose backward is a gather by token. Kimi combine 13540 -> 1102 µs (ideal 827), deterministic, 2-3x more accurate vs fp64. Under compile, Inductor's own scatter can match a gather; measure both eager and compiled before switching. ## 7. Move scaling to the narrower tensor, inside a region A linear map commutes with per-row scaling: `w2(h * s) == s * w2(h)` when w2 has no bias. Apply routing scores to the SwiGLU output (width F) inside the compiled SwiGLU, not to the w2 output (width D) in eager. DSv3 MoE module, 4k: 9.36 -> 8.17 ms (-12.8%), peak memory -29% at 16k (the saved `[R, D]` tensor goes away), lower error vs fp64. Guard it exactly: ```python use_scores_in_act = type(act) is SwiGLU and getattr(w2, "bias", None) is None and output_postprocess is None ``` In eager it is only -3% / -5% (4k / 16k) vs -12.8% / -13% with the region, so about a quarter to a third of the win survives without compile; the region is what makes it pay. Under all-to-all EP the scores live on the source rank, so this rewrite needs them sent through the all-to-all with the tokens; the gather-sum combine (pattern 6) works on every dispatcher. ## 8. Narrow sort keys `argsort` of int64 expert ids sorts all 64 key bits (22 kernels). Ids fit in int16 when the expert count does: ```python keys = expert_ids.to(torch.int16) if num_experts <= torch.iinfo(torch.int16).max else expert_ids order = torch.argsort(keys, stable=True) # identical permutation ``` 72.9 -> 30.4 µs, 11 kernels. ## 9. Unroll tiny iterative loops per entry A Sinkhorn loop on `[T, 4, 4]` realizes every iteration (each column sum reads the previous step transposed): 163 kernels, 309 µs. Writing it as 16 per-entry `[T]` tensors gives 37 kernels, ~100 µs, bitwise in fp32/fp64. Cost: much longer cold compile (+~3 min). Present it as a trade-off. ## 10. Router glue as one region Sigmoid scores, expert bias, group-limited top-k and token sums: compile them as one side-effect-free function (`_route`) and keep counters, bias updates and aux-loss injection in eager. Inductor lowers `topk` as an ATen fallback (no fusion). For group-limited routers, argmax rounds fuse instead; ungrouped routers keep `topk`. Router glue at [T, 256]: 4k 356.7 -> 135.9 µs. ## 11. Data-dependent shapes and fakes - A custom op whose output size depends on data needs a fake that returns an unbacked size (`torch.library.get_ctx().new_dynamic_size()`), not the input's size. A wrong fake makes make_fx record the wrong shape and Inductor assert at runtime. - Branches in fakes or backward on that size will guard and fail during tracing; `guard_or_true` avoids that, or better, make the real op return a shape that always matches the fake (e.g. zeros instead of an empty tensor on a rank that received nothing). - Remove host syncs from the training path: use counts the library already returned, copy from pinned memory with `non_blocking=True` when the data is final, and avoid boolean-mask indexing (sort with a sentinel and slice by a host int instead). ## 13. fp32 GEMMs Small fp32 GEMMs (DSv4 HcPre/Compressor projections, a router gate) run as slow SIMT kernels unless an emulated fp32 matmul precision (BF16x9) is on, and newer trees route fp32-output projections through dedicated modules (e.g. `HiMidLoLinear`; check what this tree has), and an upcast of a large bf16 input may be materialized first. - Compute an fp32-output GEMM from bf16 inputs with fp32 accumulation instead of upcasting the input (DSv4 HcPre: 3262 -> 1291 µs). - For a narrow fp32 backward, split the operand into exact bf16 pieces stacked along the output dim and use one bf16 GEMM with fp32 output (DSv3 router gate backward at [4096, 7168] x [256, 7168]: 0.288 -> 0.065 ms, same error vs fp64 at that K). The error depends on K: at K=16384 the piece split was ~35x worse than BF16x9, so check accuracy at your K. - Under compile, Inductor removes bf16 round trips (`x.bfloat16().float()`), so a split written as `hi = x.bfloat16(); lo = (x - hi.float()).bfloat16()` collapses. Split by bit masking instead: ```python hi = (x.view(torch.int32) & -65536).view(torch.float32) # keep the top 16 bits: exact in bf16 lo = x - hi # exact; repeat for a third piece ``` ## 14. Wrap non-traceable kernels as custom ops with exact fakes A library kernel that Dynamo can't trace causes a graph break, which inside a checkpointed block can make the whole block fall back to eager. - Register it with `torch.library.custom_op` plus a fake whose output sizes, strides and dtype match the real kernel exactly (a fake that returns contiguous outputs where the kernel returns strided ones breaks downstream ops). - DSv4 attention as a custom op: SAC block compile 302.5 -> 122.6 ms, because the block no longer graph-broke inside the checkpoint. - Check libraries that pick a different kernel under `is_compiling()`: one used its fast CuTe path eagerly and a slow Triton path under compile (8.6x slower attention). ## 15. The loss path - A loss region (final norm + lm_head + cross entropy) beat eager CE by 13% on DSv3 (4k: 34.6 -> 30.1 ms) and saved memory. - Chunked loss: one lm_head call over the concatenated chunks of all predictions instead of one call per chunk: -18% at 4k. Choose the chunk count from the token count; at large vocab, per-chunk weight-grad accumulation can cost more than full logits. ## 16. `autograd.Function` with separate forward and backward regions Compile the forward and the backward as two regions inside an `autograd.Function`, and detach inputs before calling the compiled forward. `requires_grad` is part of Dynamo's guards, so detaching the inputs lets a `no_grad` call reuse the same forward graph (no recompile), and you choose exactly what is saved for backward. ```python _fwd = local_compile("name", batch_invariant=True)(fwd_impl) # returns out and small saved stats _bwd = local_compile("name_bwd", batch_invariant=True)(bwd_impl) class Op(torch.autograd.Function): @staticmethod def forward(ctx, x, w): out, stats = _fwd(x.detach(), w.detach()) # detached: no_grad reuses this graph ctx.save_for_backward(x, w, stats) # inputs plus small stats; large intermediates are recomputed in _bwd return out @staticmethod def backward(ctx, grad_out): return _bwd(grad_out, *ctx.saved_tensors) ``` Kimi K3 attention residual: 24.3 ms eager and 12.2 ms compiled as written -> 1.56 ms per call (ideal 0.68). The same shape returns a view without a copy. A region can't return a strided view cheaply: AOTAutograd copies the view's grad into its strides (DSv3 MLA V: 556 µs), and `.contiguous()` instead costs a clone plus slower attention. Return the view from the `autograd.Function` (outside the regions) and let `_bwd` write the base's grad: DSv3 `mla_qk` 45.77 -> 45.48 ms per layer at 16k. ## 17. Move a reduction over a small symbolic dim into a GEMM A weighted sum over a small dim that varies across call sites (Kimi residual width N) lowers to a slow looped reduction when N is symbolic. Expressing it as a batched GEMM, `[T, N] x [N, D]` with the fp32 weights split into exact bf16 pieces (pattern 13), gave 180 µs vs 271 µs (static N) and 1435 µs (dynamic N). Check accuracy vs fp64. Same problem as patterns 4 and 5, different fix. ## 18. Emit an input grad in the layout the next kernel wants Inductor returns a region's input grads token-major, whatever layout the input had, so the next kernel may pay a transpose copy. Passing a transposed input alone did not change that. An identity `autograd.Function` inside the region whose backward restrides the grad did: ```python class _GradLayout(torch.autograd.Function): @staticmethod def forward(ctx, x): return x.view_as(x) @staticmethod def backward(ctx, g): return g.transpose(0, 1).contiguous().transpose(0, 1) # head-major memory, same logical shape ``` Whether it pays depends on the attention backend that consumes the grad (deceiving-quirks.md: -3.8% with Triton, neutral with FA4). ## 19. Grouped per-expert bias as a gather plus a one-hot GEMM Eager expanded each expert's bias row by row, and the backward summed bias grads with a bf16 scatter-add (2.65 ms per 2 GPT-OSS layers at 16k). Gather the bias per row in the forward, and compute the bias grad as one GEMM against a one-hot routing matrix: ```python bias_rows = bias[expert_of_row] # forward: [R, F] gathered, fuses into the GEMM epilogue # backward: one-hot [E + 1, R] (extra row for padding) @ grad_out [R, F] -> [E + 1, F], fp32 accumulation one_hot = torch.zeros(E + 1, R, device=g.device, dtype=g.dtype).scatter_(0, expert_of_row[None], 1) bias_grad = (one_hot @ g).float()[:E] ``` Deterministic, accumulated in fp32, and faster in eager too. When one region serves both w13 and w2 (two widths), mark the width static: we saw 6 graphs over the two widths and raised the limit to 16. GPT-OSS 16k: -9.7%. ## 20. Eager hooks between attention and the output projection Model-specific rescaling between attention and `wo` (GPT-OSS attention sinks) runs as a few eager kernels per layer; a region around it was -3.9% on GPT-OSS. # Writing the stack: PRs, findings doc, code Reader: a tired junior engineer. They should get each PR in one read, without opening the code. ## Writing rules - **Title:** `[][perf] `, e.g. `[dsv4][perf] Compile the mHC hyper-connections`; shared code `[moe][perf]`. - **TLDR first**, with the headline number. - **Problem and Solution as numbered baby steps**, 4-6 each: - one fact per line; each line follows from the one before; if a line needs "and", split it; - start from what the code does today, not from the fix; - end the Problem with the cost (a number); end the Solution with what stays unchanged. - **Name real things:** the function, the tensor and its shape, the config field. - **Diagram** when data moves between stages: eager vs region, fwd vs bwd, before vs after. - **Numbers in code blocks:** `before -> after`, one line per workload, conditions on one line above. - **Plain words:** explain a term on first use or say what it does ("autograd copies the two grads back into one tensor", not `UnbindBackward0`). - **Direction and both numbers:** "prevents a 2.7x slower backward (1089 vs 2987 us)", not "makes backward ~2x slower". - **Say whether real runs hit it:** which recipes enable the path. If none today, call it headroom. - **No paragraphs, no markdown tables.** One idea per line. - **No `#N` or GitHub URLs** to other people's PRs or issues. - **Last check:** can it be shorter? ## PR template ````markdown **TLDR:** **Problem** 1. 2. 3. 4. **Solution** 1. 2. 3. **Results** ```text -> ms () CUDA graphs -> ms () correctness: outputs + all grads vs eager, bf16 and fp32: <...>; vs fp64: <...> graphs: N for + no_grad (limit L) ``` You may ask: NOTE: **Test plan**

Stack summary (same in every PR)

```` ## When the template doesn't fit - **Cost step:** take the glue's own kernel time from the profile. Else "85.6 ms of the 240 ms step goes away (Results)". - **Numbers from elsewhere:** label the line: `(measured one base earlier, same code)`, `(1 process, indicative)`. - **Alternatives with numbers:** an `alternatives` section in the Results block; one line each. - **Big PR:** correctness gets its own code block. - **Region options** (`prefer_nd_tiling`, `recompile_limit`): a short list after Solution, one line each with why. - **Skill / docs PR:** TLDR, "What it contains" (numbered), "How it was validated", Test plan. - **Fix PR:** Problem / Solution as usual; Results is the correctness check, no speed number needed. - **Same commit in several stacks:** end with `Stack: 3 of 7. Same commit as /.` - **Supersedes, stack-level correctness:** FINDINGS only, not PR bodies. - **Blocked by compile:** only real blockers (a crash, a silent dtype or value change, a missing fusion worth >1%). One line in "You may ask" with the number: `Why not fold the residual add into addmm? Inductor copies the residual first: 64 extra copies per step.` None hit: say nothing. - **Stack summary:** collapsed in every PR with `<- this PR` on its line; expanded in PR 1 with Rejected, Compile limits and Not tested. ## Filled example ````markdown **TLDR:** GPT-OSS's per-expert bias backward becomes one fp32 GEMM instead of a bf16 scatter-add: 16k step -9.7%, deterministic. **Problem** 1. **Bias per expert.** Each GPT-OSS expert projection adds its expert's bias to every routed row. 2. **Expanded in eager.** `repeat_interleave` copies the bias out to all R rows, then adds it. 3. **Scatter in backward.** Autograd sums the R row grads back into the bias with a bf16 scatter-add. 4. **Cost.** The scatter alone is 2.65 ms per 2 layers at 16k tokens, and its result changes run to run. **Solution** 1. **Gather, don't expand.** Each row reads its expert's bias by index. 2. **Bias grad as a GEMM.** `add_grouped_bias` computes it as `one_hot[E+1, R] @ grad[R, O]`, accumulated in fp32. 3. **Region.** `grouped_expert_bias` fuses gather + add; width marked static (w13 and w2 differ), recompile_limit 16. 4. **Unchanged.** Expert weights, routing and the dispatcher. ```text before: bias[E,O] --repeat_interleave--> [R,O] --add--> bwd: scatter_add bf16 (atomics) after: bias[E,O] --gather(expert_id)--> add (region) bwd: one_hot[E+1,R] @ grad (fp32) ``` **Results** ```text one projection, 16 experts, R = 4T, fwd+bwd kernel time, GB300 T=16k O=5760 1593 -> 615 us (ideal ~319) T=4k O=5760 313 -> 161 us GPT-OSS 20b layers 0-1, 16k tokens, fwd+bwd wall, GB300, 1 process (indicative) step 28.46 -> 25.70 ms (-9.7%) correctness: fp64 CPU test vs the repeat_interleave reference, incl. an empty expert and padding rows graphs: grouped_expert_bias 6 (2 widths x static/dynamic rows x grad mode), limit 16 ``` You may ask: why not let block compile handle it? Inductor's scatter there uses atomics: nondeterministic. NOTE: not tested: EP>1 with padded rows from TorchAO, loss curve. **Test plan** `pytest tests/unit_tests/cpu/test_gpt_oss_compile.py` ```` ## Findings doc skeleton Its top block becomes the Stack summary of every PR. ````markdown # : local compile results ```text upstream regions -> stack 16k 240.1 -> 75.1 ms (-69%) 4k 80.2 -> 32.2 ms (-60%) baselines: eager 243.9, per-layer compile 213.1 stack (biggest first) 1 skill 2 mhc -35.7% 3 rope code 0 (enables 4) 4 attention RoPE glue -22.2% ... ``` Rejected: , one per line Compile limits hit: , one per line; omit if none Not tested: <...> Skill feedback: ```` ## Code rules - **Minimal diff.** No defensive fallbacks, shims, single-use helpers or `f -> _f -> _bound_f` chains. - **Names say what it holds;** no single letters or cryptic abbreviations. Shapes in a comment when the type is ambiguous: `routed_bias # [R, O]`. - **Docstring first line:** what the caller gets. Non-trivial public function: `Args:` plus an `Example:` with shapes; the example replaces prose. - **Comments: the local why only.** An invariant that looks unnecessary, in terms the reader has: ```python # Detached inputs: requires_grad is part of the guards, so no_grad and training share one graph. ``` No measured ratios, no narration of each line. A mechanism of 3+ steps gets numbered one-line steps. - **Config field over hardcoded value** when two users need different values (e.g. a region's `recompile_limit`). - **Gated paths side by side:** say what runs eager vs compiled in one place. - **Leave clear upstream text alone;** match the repo style guide (torchtitan's `.claude/CLAUDE.md` and `.claude/rules/`). # Symptom -> fix Read the kernel table from the profiler (see harness.md), then look up what you see. | Trace symptom | Likely cause | Try (pattern #) | |---|---|---| | `CatArrayBatchedCopy` / cat or stack kernels in eager glue, incl. a stack in the MoE backward | a cat outside a region, split-backward cats, or an eager unbind of a packed gate/up projection | put producer + cat + consumer in one region; unbind inside it; try both cat lowerings (2, 5) | | `triton_red_*` / `triton_per_*` over a dim of size 2-16, or a reduction with `XBLOCK=64` far off ideal | small reduction not unrolled; the dim may be symbolic | static: Python loop, `unroll_reductions_threshold`, `mark_static` inside the region (4, 5); varying width: `autograd.Function` or a GEMM (4, 17) | | fp32 SIMT GEMM (`sgemm`, cutlass simt), or a large fp32 copy before a GEMM | fp32 matmul precision differs from the trainer's, or an upcast materialized | match the trainer's fp32 matmul precision, then 13 | | Region kernel with many masked loads, slow vs ideal around a cat | pointwise-cat lowering on a bad shape | `max_pointwise_cat_inputs=0` (2) | | `gemv`/`gemvx` or `extern_kernels.bmm` with a tiny inner dim, plus a separate fp32 copy | extern GEMV | explicit weighted sum (3) | | `elementwise_kernel<...BinaryFunctor...>` non-vectorized broadcast mul on a wide tensor | per-row scale in eager on the widest tensor | move the scale to a narrower tensor inside a region (7) | | `aten` complex mul / `view_as_complex` kernels surviving under compile | no complex codegen | real arithmetic (1) | | `index_put`, `scatter_add`, `indexing_backward_kernel`, atomics | scatter combine or gather backward | gather-sum `autograd.Function` (6); measure Inductor's compiled scatter first | | Same region compiled twice: static then dynamic graph | a small dim became symbolic | `mark_static` on fixed dims, `maybe_mark_dynamic` on T in eager (5) | | `recompiles` log: guard on `4096 <= size` | mix-order heuristic | `triton.mix_order_reduction_non_strict_mode` | | `recompiles` log: `GLOBAL_STATE changed: grad_mode` | no_grad pass | budget; fewer call signatures | | `FailOnRecompileLimitHit` | too many signatures for one code object | fewer layouts per region, non-strict mix-order, split regions | | `cudaStreamSynchronize`, DtoH memcpy, CPU gaps between kernels | `.item()`, `.tolist()`, boolean-mask indexing, blocking copies | use returned counts, `non_blocking=True` from pinned memory, sentinel sort + host-int slice (11) | | Many tiny kernels from a loop over a tiny matrix (`[T, 4, 4]`) | each iteration realized | per-entry unrolling (9) | | Library kernel runs differently under compile (e.g. Triton instead of CuTe) | library branches on `is_compiling()` | check the library; wrap its fast path as a custom op (14) | | Graph break at a library call | non-traceable kernel | custom op with a correct fake (14) | | Region slower at small T only | host time per region call | merge regions, or CUDA-graph the whole model over the regions | | `direct_copy_kernel` on a region-output-shaped tensor in the backward | the region returns a strided view; its gradient arrives in another layout and AOTAutograd inserts a copy | return a dense (contiguous) tensor under compile | | Kernel over `[T, N, D]` with a symbolic small N, index math dividing by N, far off ideal | 1-D tiling folds the symbolic N into every index | `triton.prefer_nd_tiling=True` on the region (gating-and-options.md) | --- name: local-compile-perf description: Make a torchtitan model train faster with `local_compile` regions and compile-friendly rewrites: measure, find the slow glue, rewrite it, verify. Use when asked to speed up a model, layer or op with torch.compile, or to review such a change. --- # Speeding up a model with local compile The method: measure a model's layers against their ideal time, find the glue ops that run far slower than the bytes they move, and rewrite or regroup them so Inductor fuses them into fast kernels. Then verify numerics, graph counts and parallelism, and deliver small changes. Terms: - **Region:** a function decorated with `@local_compile` (`torchtitan/distributed/local_compile.py`), compiled on its own with `torch.compile(fullgraph=True)` when listed in the model's `local_compile_regions`. It traces through everything it calls (callees are inlined); `fullgraph` only forbids graph breaks. - **Glue:** everything that isn't a GEMM or attention: norms, RoPE, residual adds, routing, permutes, casts, scores. - **Ideal time:** minimum bytes moved / measured HBM bandwidth (or a FLOP floor for GEMMs). - **Per-layer (block) compile:** `torch.compile` on each transformer block. Used here as a baseline, not a deliverable. - **Graph budget:** compiled graphs per region; `recompile_limit` is 8 by default and `fullgraph` turns overflow into an error. - **Gate:** an if-branch that picks the rewritten form only when compiling. - **fp32 matmul precision:** e.g. BF16x9, fp32 matmuls emulated on bf16 tensor cores (`torch.backends.cuda.matmul.fp32_precision`). Trainers differ across trees; match what this tree's trainer sets. - **SAC:** selective activation checkpointing. **GraphTrainer:** torchtitan's whole-step tracer (make_fx + Inductor). ## Mental model ```text eager: x -> [norm] -> HBM -> [rope] -> HBM -> [cat] -> HBM -> out (3 kernels, 3 round trips) region: x -> [norm + rope + cat] -> out (1 kernel, inputs read once) win available = measured glue time - ideal time ``` - GEMMs and attention dominate a step; compile rarely speeds them up. The budget is the glue. - Compile wins when Inductor can see and fuse the ops. Complex numbers, opaque custom ops, extern GEMV, data-dependent shapes and symbolic small sizes block or degrade fusion; rewriting around them is most of the work. - Every region costs graphs, cold-compile seconds and host launch time. At small token counts, host time can cancel kernel wins; CUDA-graphing the whole model after the regions are applied removes that host time on replay, which is why every number is also reported under CUDA graphs (step 7). ## Sizing a region Aim for the smallest region that captures the win: ```text C() = A() + B() only A wins compiled -> region(A) A and B win independently -> region(A), region(B) fused kernels span A|B and beat A + B compiled apart -> region(C) (C inherits all call signatures -> more graphs) ``` Example: DSv4 q norm + RoPE went from 1497 µs eager to 1044 µs in one region, because the cat fused into the norm's reduction. ## Workflow ```text 0 plan (devlog, stop rule) 1 baselines ............. harness.md 2 ideal times ........... measurement.md 3 profile + diff ........ harness.md (profiler) -> symptoms.md 4 rank candidates 5 pick a fix ............ patterns.md, gating-and-options.md 6 graph budget .......... measurement.md#graph-budget 7 screen, then confirm .. measurement.md 8 correctness, AC, TP/EP measurement.md, gating-and-options.md, deceiving-quirks.md 9 commit + scope 10 stack + findings doc . Output (below), pr-template.md 11 feed the skill back Before investing in an idea: negative-results.md. If compile itself fails: compile-limits.md. ``` 0. **Plan.** Set a time budget and a stop rule: stop when the remaining gap to ideal is under ~1% of the layer, or the next candidate's estimated gain is below that. Keep a devlog: one file per effort, a timestamped entry appended right after each measurement, decision and dead end (not in batches). Check how this tree declares regions (harness.md: where regions live) and what fused kernels already exist under `torchtitan_recipes/overrides/` (e.g. `fused_mla`); benchmark them as alternatives. 1. **Baselines** (harness.md): - 2 layers at real dims, covering each block type; packed tokens per rank from small to large (e.g. 2k, 4k, 16k, 32k+). Torchtitan has no batch dim, so long rows are normal. - fwd+bwd, with the recipe's activation checkpointing mode (none, full, selective). - three trees: eager, upstream regions (what the model's `local_compile_regions` enables on main today), per-layer compile. Rewrites plus regions beat per-layer compile by 18-29% on Kimi K3, so it's a baseline, not a ceiling. - fp32 matmul precision matching this tree's trainer. - check that the expected attention backend actually loaded (kernel names in the profile) and pin optional deps (requirements.txt); a missing FA4 or a stale attn-gym silently falls back to a slow path (deceiving-quirks.md). - record wall and kernel time, peak memory, kernel count; check attention against its FLOP floor. When attention is the largest item and far from its floor, it's a backend choice (flex, varlen, FA4, cuDNN): report it with the floor rather than trying a region. - per-layer compile caveats: attention inlined into an outer compile loses its own region options (flex attention got 25% slower), so compare glue only; on an older tree, `torch.compile` on a SAC `CheckpointWrapper` silently ran eager kernels (not reproduced on current main; check the profile). 2. **Ideal times and best known kernels** per glue op; build a measured / ideal / best-known table. The best known kernel for the same op (the tree's overrides, quack, Liger, Megatron, Transformer Engine, FA4) is a reachable target, and diffing its profile against yours kernel by kernel shows where the gap is. DSv3: FusedMLA proved 1.24 ms was reachable in `mla_qk`; the diff found RoPE's complex mul (0.42 ms). 3. **Profile and diff** eager vs regions vs per-layer compile by kernel; look up each symptom in symptoms.md. 4. **Rank** by (measured − ideal) × share of the layer. Items under ~1% are usually not worth a PR unless they fix correctness or remove a host sync. 5. **Pick a fix**, roughly cheapest first: ```text enable an existing region resize a region (A / B / C above) rewrite the function so it lowers well (patterns.md) add a region-scoped option (gating-and-options.md) autograd.Function with a hand-written backward ``` Fixes interact: before rejecting a rewrite or an option, try it with the others. DSv3 `mla_qk`: real-arithmetic RoPE alone was slower (46.95 ms), the cat option alone neutral (46.37), both together 46.34 -> 45.85 ms. 6. **Graph budget, early:** count all shapes and call sites plus a `no_grad` pass. Levers: `mark_static`, a dynamic token dim marked in eager, fewer signatures per region, `triton.mix_order_reduction_non_strict_mode`, raising the limit (we used 16 in several models; how depends on the tree, see measurement.md). Details in measurement.md. 7. **Screen in one process, confirm with 3 fresh processes** on a quiet GPU. Every comparison reports two rows: normal launch and CUDA-graph replay of the whole fwd+bwd (harness.md). A win that disappears under graphs was host time. 8. **Correctness and side effects:** measurement.md (correctness) and gating-and-options.md (numerics, AC, parallelism); check deceiving-quirks.md. 9. **Commit each validated change locally** before the next. Enable new regions in the target model's config or recipe, not a shared default list. A region enabled in a model config also reaches every config that subclasses it (DSv3 -> Kimi K2.7): check subclasses or enable per recipe. List which recipes reach the change. 10. **Package** the stack and findings doc (Output). 11. **Feed the skill back:** list what this skill missed (new patterns, edge cases, negative results, wrong advice, new compile limits for compile-limits.md) as proposed additions. ## Output - **A diff stack:** one small branch per change, each on the previous one, committed locally, ordered by impact (biggest first; state the order). Ask the user whether to open it as draft PRs, and where, or not at all. Don't open PRs without asking, and don't reference other people's PRs or issues. - **A findings doc** (short): what was measured, the stack with deltas, rejected ideas, not tested, skill feedback. Its top block is the stack summary in every PR. - **PR text:** TLDR, Problem and Solution as numbered baby steps, a diagram, numbers in code blocks; no paragraphs or tables. Template, filled example and rules: pr-template.md. - **Code:** minimal diff; regions enabled per model config; a test that fails before and passes after; code rules in pr-template.md. - **Compute:** - iterate on one GPU; - multi-GPU only for what one GPU can't show (comm overlap, EP/TP behavior, zero-token ranks): build-check the harness at the target world size (meta tensors or a fake process group), then submit one small, short job and keep working; - fake process groups and FakeTensor for tracing and shape checks only; above world 1, fake peers alias storage and give garbage routing; - parallel subagents for independent work; a fresh non-fork agent for blind checks. ## What usually saves time Lessons from specific cases, not rules: - **Measure the function as written under compile before rewriting.** Inductor sometimes already does well: a compiled scatter matched our gather rewrite. - **Move work to the narrower tensor.** Applying routing scores to the 2048-wide SwiGLU output inside its region, instead of the 7168-wide w2 output in eager, saved 13% of a MoE layer and 29% peak memory. - More in negative-results.md. ## Checklist - [ ] Plan: budget, stop rule, devlog. - [ ] Baselines (eager / upstream / per-layer) at 3+ token counts, recipe AC mode, trainer's fp32 matmul precision, with and without CUDA graphs. - [ ] Ideal-time table; kernel diff; ranked candidates (ms, % of layer). - [ ] Region boundaries chosen by measurement. - [ ] Graph count: all shapes + `no_grad`, headroom under the limit. - [ ] 3 fresh processes; outputs and all grads vs eager (bf16, fp32) and fp64. - [ ] Gate chosen from eager cost (gating-and-options.md). - [ ] Quirks checked: TP, EP, hardware, dispatch mode, AC, compile mode. - [ ] Regions per model config; reaching recipes listed; each change committed. - [ ] Stack + findings doc; asked before opening any PR; skill feedback. # Customizing the capture and diff ## 1. Excluded ops (`_EXCLUDED_OPS` in `activation_tracer.py`) Denylist of op names dropped at capture time as "infrastructure, not numerics." Includes view/reshape/permutation/indexing/creation ops, type-casts (`_to_copy`, `to`, `clone`), and `detach`. - **Add to the list** if a new op is firing that pollutes the diff with structurally identical rows (e.g. a fused kernel that wraps things you already capture upstream). - **Remove from the list** if you want visibility into something you currently hide. Collective comms (`all_gather_into_tensor`, `reduce_scatter_tensor`, `wait_tensor`, etc.) are *not* in the default list — they show up in the diff because they often diverge between eager and traced. Op names skipped by the filter that *actually fired* during capture are surfaced as a chip list per side in the HTML. Use that to discover what you're dropping. ## 2. Numel / dtype filter (`_should_capture` / `min_numel`) By default we skip tensors smaller than `min_numel=1000` and tensors that aren't `{float32, float16, bfloat16}`. Adjustable via the `min_numel` and `op_filter` kwargs on `DebugModeTracer` (constructed directly; not a CLI flag). ## 3. Hash function (`log_tensor_hashes`) The `output_hash` / `input_hashes` fields come from `DebugMode.log_tensor_hashes(hash_fn="norm", hash_inputs=True)`. `"norm"` = L1 in float64. Other options DebugMode supports: - `"hash_tensor"` — `torch.hash_tensor` (XOR-reduce of the underlying bytes; truly bit-exact, fast, but loses arithmetic-tolerance semantics). - A custom callable — `lambda t: (t.float().mean().item(), t.float().std().item())` for mean+std as a "hash." To switch, edit the `DebugMode.log_tensor_hashes(hash_fn=..., hash_inputs=True)` call inside `DebugModeTracer.__enter__`. ## 4. Manual overrides (`--override `) The four-pass matcher (override / exact key / fuzzy key / stats) handles most cross-mode pairings automatically, but some patterns don't match: - **AC-recompute attribution drift**: eager's selective-activation-checkpoint recompute fires modules directly, so ModTracker only sees `feed_forward` (no layer prefix), while traced sees the full `layers.N.feed_forward`. - **Per-layer counter shifts**: an in-place collective op (`_allgather_base_`) takes the `op_0` slot in eager's per-layer counter; traced doesn't have it under the layer key, so subsequent ops shift by one. - **Collective renaming**: eager's in-place `_allgather_base_` ↔ traced's lowered `all_gather_into_tensor_out` + `wait_tensor` chain. Override file format (`# `-prefixed lines are comments): ``` # AC-recompute feed_forward/op_7_mul, layers.2.feed_forward/op_3_mul # Counter shift layers.1/op_1_add, layers.1/op_0_add # Collective rename layers.0/op_0__allgather_base_, layers.0.attention_norm/op_1_all_gather_into_tensor_out ``` Each line force-pairs `, ` regardless of stat similarity. ### Common eager-vs-traced override scenarios Confirm each candidate with matching `Shape` and matching or near-matching `output_hash` before adding it. - **AC recompute FQN drift**: eager may log recompute under a short FQN, while traced keeps the full layer FQN. `feed_forward/op_7_mul, layers.2.feed_forward/op_3_mul` - **Per-layer counter shifts**: eager FSDP collectives can consume early `op_N` slots, so later same-layer compute has shifted counters. `layers.0/op_2_add, layers.0/op_0_add` - **Repeated attention math**: Qwen-style rotary/GQA ops repeat similar `mul` / `neg` / `add` / `bmm` patterns, so fuzzy matching can drift by one repeated block. `layers.0.attention/op_0_mul, layers.0.attention/op_8_mul` - **Backward attribution drift**: backward ops can move to ``, a parent module, or a generated backward op name. `layers.0.attention/op_12_add, /op_66_add` - **Communication lowering**: eager `_allgather_base_` / `_reduce_scatter_base_` often corresponds to traced `wait_tensor` rows under the consumer module. `layers.0/op_1__allgather_base_, layers.0.attention_norm/op_1_wait_tensor` ## 5. HTML appearance - `--name1` / `--name2` set the run labels in headings and the side column. - Cell coloring is log-scale on relative diff (`1e-8` → faint pink, `1e-1` → full red). Adjust the thresholds in `_hash_cell` and the stat-cell loop in `generate_html` if you want a different gradient. - Input-hash cells are collapsible (3-line max by default). Adjust `.ih-cell.collapsed { max-height: 3em }` in the CSS for a different cap. # Patching torchtitan to enable capture The capture is gated on a profiler flag and an `ActivationCaptureProfiler` constructed inside `Profiler`. Apply these patches before a capture run, then revert them when you are done (they don't belong on `main`). ## 0. Import bootstrap `activation_tracer.py` lives in `.claude/skills/numerics_debugging/scripts/`, which is not an importable package (`.claude` is not a valid Python identifier), so every patch site below has to put that directory on `sys.path` before importing from it: ```python def _numerics_scripts_on_path() -> None: """Put the numerics_debugging skill's scripts/ on sys.path.""" import sys from pathlib import Path for parent in Path(__file__).resolve().parents: scripts = parent / ".claude" / "skills" / "numerics_debugging" / "scripts" if scripts.is_dir(): if str(scripts) not in sys.path: sys.path.insert(0, str(scripts)) return raise RuntimeError("numerics_debugging skill scripts/ not found") ``` Each snippet below calls this immediately before its import. Paste the helper into whichever file needs it, or into one shared spot and import it. Walking up from `__file__` matters: it anchors on the checkout that contains the file you are patching. Anchoring on `torchtitan.__file__` instead looks equivalent but is not -- with an editable install, `import torchtitan` can resolve to a *different* checkout depending on the working directory, and the bootstrap would then load another tree's tracer or fail outright. ## 1. `torchtitan/observability/profiler.py` Add the import, config field, constructor arg, lifecycle hooks, and builder. ```python # top of file _numerics_scripts_on_path() from activation_tracer import ActivationCaptureProfiler # inside Profiler.Config — add next to enable_memory_snapshot dump_numerics: bool = False """Dump per-op activation logs for numerics debugging. Writes ``{dump_folder}/numerics/rank_{rank}_activations.log`` (per-op stats + norm hashes of inputs / outputs).""" # Profiler.__init__ — add model kwarg and slots def __init__( self, config: "Profiler.Config", *, global_step: int = 0, base_folder: str = "", leaf_folder: str = "", model: torch.nn.Module | None = None, # for activation capture profiler ) -> None: ... self.activation_capture_profiler = None # ActivationCaptureProfiler registers global module forward hooks on # the model so backward ops can recover their owning FQN. self._model = model # Profiler.__enter__ — build the activation capture profiler alongside the memory profiler self.activation_capture_profiler = self.build_activation_capture_profiler( base_folder=self._base_folder, ) # Profiler.__exit__ — teardown if self.activation_capture_profiler is not None: self.activation_capture_profiler.__exit__(exc_type, exc_val, exc_tb) self.activation_capture_profiler = None # Profiler.step — drive the capture-step cadence if self.activation_capture_profiler is not None: self.activation_capture_profiler.step() # new method def build_activation_capture_profiler(self, *, base_folder: str): """Create and return an :class:`ActivationCaptureProfiler`, or ``None`` if disabled.""" cfg = self._config if not cfg.dump_numerics or self._model is None: return None dump_dir = os.path.join(base_folder, "numerics") profiler = ActivationCaptureProfiler( enabled=True, model=self._model, dump_dir=dump_dir, capture_step=cfg.profile_freq, ) profiler.__enter__() return profiler ``` ## 2. `torchtitan/trainer.py` Pass the model to the profiler when entering the training loop: ```python with config.profiler.build( global_step=self.step, base_folder=config.dump_folder, model=self.model_parts[0], # add this line ) as profiler: ... ``` `model_parts[0]` is the eager model rank-0 owns; that is what `ActivationCaptureProfiler` installs forward hooks on so DebugMode's ModTracker can attribute backward ops to the right FQN. ## 3. graph_trainer only — replay traced graph through FQNInterpreter When the active recipe sets `compile.mode = "aot_fx_trace"`, the traced graph is called as `gm(*flat_inputs)`, which bypasses every `nn.Module.forward`. That means `DebugMode`'s `ModTracker` can no longer attribute ops to a FQN, and the log degrades to `/op_N_*` everywhere. The fix is to walk the graph node-by-node with an FX interpreter that restores the FQN / stack / phase that the traced commit already stashed in `node.meta`, so capture gets the same context that eager would. **3a. `torchtitan/experiments/graph_trainer/debug_utils.py`** — append: ```python class FQNInterpreter(torch.fx.Interpreter): """Interpreter that sets activation tracer context vars from node metadata. For each node, reads: - ``node.meta["custom"]["module_fqn"]`` → sets _current_module_name - ``node.meta["stack_trace"]`` → parsed and set as _current_stack_frames - ``node.meta["autograd_backward"]`` → sets _current_phase_override This is needed because traced graph replay via ``gm(*inputs)`` bypasses module forwards entirely, so DebugMode's ModTracker cannot infer the module FQN. FQNInterpreter walks the graph node-by-node, restoring the context that eager capture would otherwise get from DebugMode and autograd metadata. """ def run_node(self, n: torch.fx.Node): from contextvars import Token _numerics_scripts_on_path() from activation_tracer import ( _current_module_name, _current_phase_override, _current_stack_frames, _parse_stack_trace, ) fqn = (n.meta.get("custom") or {}).get("module_fqn") stack_trace = n.meta.get("stack_trace") is_backward = n.meta.get("autograd_backward", False) phase = "backward" if is_backward else "forward" tokens: list[Token] = [] if fqn: tokens.append(_current_module_name.set(fqn)) if stack_trace: tokens.append(_current_stack_frames.set(_parse_stack_trace(stack_trace))) tokens.append(_current_phase_override.set(phase)) try: return super().run_node(n) finally: for token in reversed(tokens): token.var.reset(token) ``` `run_traced` already accepts an `interpreter_cls` kwarg (in `torchtitan/experiments/graph_trainer/make_fx_tracer.py`), so no patch is needed there — just pass `FQNInterpreter` through. **3b. `torchtitan/experiments/graph_trainer/trainer.py`** — pick the interpreter only when capture is active, and thread it through `run_traced`: ```python def _maybe_get_fqn_interpreter(self) -> type | None: _numerics_scripts_on_path() from activation_tracer import ( is_numerics_capture_active, ) if is_numerics_capture_active(): from torchtitan.experiments.graph_trainer.debug_utils import FQNInterpreter return FQNInterpreter return None # in forward_backward_step, where run_traced is invoked: outputs = run_traced( ..., interpreter_cls=self._maybe_get_fqn_interpreter(), ) ``` The interpreter only kicks in on the capture step (and only under aot_fx_trace), so steady-state training is untouched. --- name: numerics_debugging description: Capture and compare per-op activations between two TorchTitan runs to spot numerics divergence (eager vs aot_fx_trace, FSDP vs no-FSDP, before vs after a refactor). Use when the user wants to debug bitwise / numeric drift in training, or invokes /numerics_debugging. --- # Numerics Debugging (DebugMode-based) Per-op activation capture + comparison toolkit. Captures activations on a designated step via `torch.utils._debug_mode.DebugMode`, then diffs two captures into an HTML report to surface numerics divergence between runs that should agree (bitwise, or within float32 reduction-order noise). Two pieces live in `scripts/`: - `activation_tracer.py` — runtime capture, driven from `Profiler` by `ActivationCaptureProfiler`. Output: `{dump_folder}/numerics/rank_{N}_activations.log`. - `compare_numerics.py` — diffs two logs, produces an HTML report. Standard library only, no torch import. They sit **outside** the torchtitan package on purpose. Nothing in core torchtitan or `graph_trainer` references them; an agent must edit torchtitan to wire the tracer in before a capture run, and revert the edits when done (they don't belong on `main`). Because `.claude` is not a valid Python package name, `activation_tracer` is imported by putting `scripts/` on `sys.path` rather than by dotted module path — see [references/patching.md](references/patching.md). > The two runs being compared **must use the same dtype and seed**. The > matcher keys on shape + float64 L1 norm; a precision change (bf16 vs > fp32) makes every row diverge and the matcher degrades to the > structural-only `stats` pass. ## Workflow 1. **Patch torchtitan** to wire the capture into `Profiler` (and `graph_trainer` if you're capturing the traced path). Full patch set: [references/patching.md](references/patching.md). 2. **Define capture recipes** for the two runs. Each recipe should derive from the configuration under test and set the capture controls in Python: ```python def numerics_run_a() -> Trainer.Config: config = base_config_a() config.training.steps = 2 config.profiler.dump_numerics = True config.profiler.profile_freq = 2 config.debug.seed = 42 config.debug.deterministic = True config.training.mixed_precision_param = "float32" return config ``` Then capture once per recipe: ```bash MODULE=my_numerics_configs CONFIG=numerics_run_a ./run_train.sh \ --output-dir ./outputs/run_A ``` Repeat with `numerics_run_b` and `./outputs/run_B`. The capture step is `profile_freq`. With `profile_freq=2` and `training.steps=2`, step 1 warms up and step 2 is the snapshot. Capture adds ~10–40% memory only on the capture step (stats are computed inline in float64; tensors aren't held). 3. **Diff** the two logs: ```bash python .claude/skills/numerics_debugging/scripts/compare_numerics.py \ outputs/run_A/numerics/rank_0_activations.log \ outputs/run_B/numerics/rank_0_activations.log \ --name1 run_A --name2 run_B \ -o diff.html ``` Open `diff.html`. Each row pairs one op from each run; cells turn red when a stat diverges; the "Match method" chip shows which of the four matching passes (override / exact key / fuzzy key / stats) paired the row. ## Customizing Excluded ops, numel / dtype filter, hash function, the manual-override file format, and HTML appearance are all tunable. See [references/customization.md](references/customization.md), which also catalogs the common eager-vs-traced mismatch patterns (AC-recompute FQN drift, per-layer counter shifts, collective renaming) you'll see in the diff and how to express them as overrides. --- name: torch_bisect description: Bisect PyTorch commits to find the regression that breaks TorchTitan. Use when the user wants to bisect PyTorch or invokes /torch_bisect. disable-model-invocation: true --- # PyTorch Git Bisect for TorchTitan Bisect the PyTorch repo to find the commit that broke TorchTitan. Always use absolute paths — shell state does not persist between Bash calls. ## Phase 1: Gather Information Record the TorchTitan directory (`pwd`). Then check if `~/.claude/torchtitan_torch_bisect_cache.json` exists — if so, show cached values and let the user reuse or override them. Collect via AskUserQuestion (pre-fill from cache if available): 1. **PyTorch repo path** — absolute path to local checkout 2. **Build command** — shell command to recompile PyTorch after each checkout 3. **Test command** — run from TorchTitan dir; must exit 0 on good commits, non-zero on bad. Judge ONLY by exit code, never by log content. 4. **Good commit** — a commit hash or a date (e.g., `Feb 15`) 5. **Timeout policy** — if test exceeds 10 min, auto-mark bad or ask each time? After collecting, offer to save answers to the cache file. ## Phase 2: Setup 1. **Check GitHub connectivity**: run `curl -s -o /dev/null --connect-timeout 20 -w "%{http_code}" https://github.com`. If it fails or returns non-200, tell the user GitHub is not reachable and suggest checking internet connectivity or proxy settings. On Meta internal servers, suggest setting `https_proxy=http://fwdproxy:8080`. Do not proceed until connectivity is confirmed. 2. `cd && git fetch origin` (timeout 300000ms). On failure, ask user to fix (remind about proxy/fwdproxy if it looks like a network error). 3. Check `git bisect log 2>&1` — if a bisect is active, ask user to reset or abort. 4. Resolve the good commit: - Hash: verify with `git log -1 ` - Date: `git log origin/main --before="T23:59:59" --format="%H %s" -1` — show result and confirm with user. 5. Bad commit is `origin/main`. Show it to user. 6. Show `git rev-list --count ..origin/main` and estimated steps (~log2). 7. Run each bisect command individually so failures can be reported cleanly: - `git bisect start` - `git bisect bad origin/main` - `git bisect good ` If any step fails, show the error and run `git bisect reset` before asking the user how to proceed. ## Phase 3: Bisect Loop Repeat until bisect finds the first bad commit: **A. Build** — run build command in PyTorch dir (timeout 600000ms). On failure: show last 50 lines, ask user to **Skip** / **Retry** / **Abort**. **B. Test** — run test command in TorchTitan dir (timeout 600000ms). Exit 0 = good, non-zero = bad. On timeout, follow user's policy from Phase 1. Never analyze logs — use only exit code. **C. Record** — `git bisect good` or `git bisect bad` in PyTorch dir. - Output contains "is the first bad commit" → done, go to Phase 4. - Output contains "Bisecting:" → continue loop. - Cannot narrow further (too many skips) → tell user, suggest abort. **D. Progress** — print step number, commit hash, good/bad, remaining count. ## Phase 4: Report and Cleanup 1. Capture `git bisect log` and `git show --stat `. 2. Extract PR number from commit subject (`(#NNNNN)` pattern). If found, run `gh pr view --repo pytorch/pytorch --json title,body,url` to get PR details. Fallback URL: `https://github.com/pytorch/pytorch/pull/`. 3. `git bisect reset`. 4. Present summary: ``` ## Bisect Complete **First bad commit:** **Commit message:** **Author:** | **Date:** ### Associated Pull Request **PR:** # — **Link:** https://github.com/pytorch/pytorch/pull/<N> **Summary:** <first ~20 lines of PR body> ### Changed files <git show --stat output> ### Full bisect log <git bisect log output> ``` Omit the PR section if no PR number was found. Do NOT diagnose root cause. ## Rules - Always use absolute paths. - Judge test results by exit code only — never analyze output. - On unexpected errors, show the error and ask the user before proceeding. - Keep the user informed after every bisect step.  --- name: add_moe_model description: Add a new HF MoE model to the transformers_modeling_backend experiment. Use when the user provides a HuggingFace model ID and wants to integrate it with Titan's MoE, or invokes /add_moe_model. --- # Add HF MoE Model to Transformers Backend Add support for a new HuggingFace MoE model in the `torchtitan/experiments/transformers_modeling_backend/` experiment. ## Scope of changes **ONLY modify code under `torchtitan/experiments/transformers_modeling_backend/`.** If supporting the new model appears to require a change outside this folder (e.g. in `torchtitan/models/`, `torchtitan/distributed/`, `torchtitan/components/`, or any other core path), **STOP and escalate to the user** — describe the change you believe is needed and why, and wait for their decision instead of editing core yourself. The one exception is the **temporary, must-be-reverted diagnostic patches** in Phase 3c (numerical-gap analysis): those edit core only to measure where titan diverges from HF, are reverted before the change is done, and are never committed. Any *permanent* change outside the experiment folder still requires escalation. **DO modify experiment code** to support the new model. The experiment folder (`torchtitan/experiments/transformers_modeling_backend/`) is designed to be extended for new models. Expected changes include: - `state_dict_adapter.py` — add regex patterns for new key naming conventions, handle transposed weight layouts, fix round-trip mappings - `moe_replacement.py` — add detection logic for new scoring functions, router types, expert layouts, or shared expert patterns in `_probe_hf_moe_block` and its helper functions - `hf_sharding.py` — give **every parameter/buffer-bearing module** the model introduces a sharding config. This is not limited to attention projections: `set_hf_sharding_configs` covers the root-level `tok_embeddings`, `norm`, `lm_head`, and `rotary_emb` (and their buffers), and `_set_layer_sharding_configs` covers per-layer norms, the attention block, and the dense MLP. A module with params/buffers but no config will mix a plain tensor with a DTensor and crash under TP (the `_assert_all_states_sharded` backstop now makes that a loud setup-time error for any param/buffer-bearing dense-path module — root modules and every decoder layer, excluding the Titan-swapped MoE). For a genuinely new module, shard it (`colwise_sharding_config`/`rowwise_sharding_config`) if it has a clear TP layout, otherwise default to `_replicate_sharding_config(module)` (replicate weights, no TP). If a new module type doesn't fit any existing helper, that's a signal to stop and think (or escalate) rather than leave it unconfigured. - `model.py` — handle alternative submodule names: HF models name the same component differently (feed-forward `mlp` vs `feed_forward`, final norm `norm` vs `final_layernorm`). Add a new model's variant to the relevant lookup (`_get_moe_attr_name`, the `HFTransformerModel` accessor properties). - `scripts/numerical_equivalence.py` — add the model's synthetic test config When blockers are found in Phase 1, fix them in the experiment code before proceeding to Phase 2. Do NOT stop and ask — implement the fixes, then continue the workflow. ## Input The user provides a HuggingFace model ID (e.g. `Qwen/Qwen3-30B-A3B`). Optionally, a local path to downloaded weights for pretrained testing. Before starting, ask the user (if not already specified): 1. **Which branch to work on** — should changes go on the current branch, or a new branch? If running in a worktree, clarify whether the commit should target the main feature branch or stay on the worktree branch. 2. **Whether to commit** — after all tests pass, should the changes be committed automatically, or left as uncommitted changes for the user to review first? ## Phase 1: Compatibility Probe Load the model config and a single MoE layer on meta device to verify titan compatibility. No file changes in this phase — but record all issues found. After the probe, fix any blockers in the experiment code before moving to Phase 2. ### 1a. Load and inspect config ```python from transformers import AutoConfig, AutoModelForCausalLM config = AutoConfig.from_pretrained(model_id, trust_remote_code=True) ``` Record: `model_type`, `hidden_size`, `intermediate_size`, `moe_intermediate_size`, `num_attention_heads`, `num_key_value_heads`, `num_hidden_layers`. ### 1b. Instantiate on meta device ```python with torch.device("meta"): model = AutoModelForCausalLM.from_config(config, trust_remote_code=True) ``` ### 1c. Find MoE layers Scan `model.model.layers` for layers where `hasattr(mlp, "gate") or hasattr(mlp, "router")` AND `hasattr(mlp, "experts")`. Record: - Which layer indices are MoE vs dense - The mechanism: `first_k_dense_replace`, `decoder_sparse_step`, or `mlp_layer_types` - Minimum `num_hidden_layers` needed to have at least one MoE layer ### 1d. Probe MoE block Call `_probe_hf_moe_block(moe_block, config)` from `torchtitan/experiments/transformers_modeling_backend/moe_replacement.py`. **Check each field succeeds without error:** | Field | How it's resolved | Blocker if missing | |-------|-------------------|--------------------| | `num_experts` | `experts.num_experts`, `gate.weight.shape[0]`, or config attrs | Yes | | `dim` | `config.hidden_size` | Yes | | `moe_intermediate_size` | `experts.gate_up_proj.shape[1] // 2` or config | Yes | | `top_k` | `moe_block.top_k`, `config.num_experts_per_tok` | Yes | | `score_func` | `sigmoid` if `e_score_correction_bias` exists, else `softmax` | Blocker if neither | | `route_norm` | `config.norm_topk_prob` or default `True` | No | | `route_scale` | `config.routed_scaling_factor` or `1.0` | No | | `num_expert_groups` | `config.n_group` or `None` | No | | `shared_expert_info` | Probes `shared_expert`/`shared_experts`/`shared_mlp` | No | **Issues to detect and fix in experiment code:** - Expert weight layout transposed (e.g. `(E, H, 2*I)` instead of `(E, 2*I, H)`) — fix `state_dict_adapter.py` to transpose during conversion, and `moe_replacement.py` to detect the layout - Experts stored as `nn.ModuleList` of individual MLPs instead of fused 3D tensors — this is a true blocker (needs HF upstream change), report to user and stop - Scoring function not `softmax`/`sigmoid` but detectable from the model code (e.g. `torch.sigmoid` call without `e_score_correction_bias`) — fix `_resolve_score_func()` in `moe_replacement.py` - Router uses `nn.Linear` directly (key is `router.weight` not `gate.weight`) — add pattern to `state_dict_adapter.py` - MoE layer attribute is not `mlp` (e.g. `feed_forward`) — fix `model.py` and `moe_replacement.py` to handle the alternative name - Shared expert round-trip fails (singular/plural mismatch) — fix reverse patterns in `state_dict_adapter.py` - Model doesn't support `grouped_mm` expert implementation — adjust the numerical test to skip the `grouped_mm` alignment for this model - Layer-level MoE: router and experts are siblings of the dense MLP at the decoder layer level, not nested inside `layer.mlp` (e.g. Gemma4 has `layer.mlp` + `layer.router` + `layer.experts`). In this case the dense MLP is functionally a **shared expert** — it runs on every token and its output is summed with the routed expert output. The correct mapping is to treat the dense MLP as `shared_experts` inside Titan's `MoE` module and replace `layer.mlp` with the full Titan MoE (router + experts + shared_experts). This way the entire block has a single `ShardingConfig` and TP/EP work correctly. Fix in `moe_replacement.py`: when probing, detect that the dense MLP is a shared expert, include it in `shared_expert_info`, and when swapping, replace `layer.mlp` with the Titan MoE that includes the dense MLP as its shared expert. Wire the router and expert weights from the layer-level attributes into the Titan MoE. ### 1e. Probe attention for ShardingConfig Check which sharding configs `_set_layer_sharding_configs` (in `hf_sharding.py`) needs: Q and KV are detected **independently**, so any mix of full-rank/low-rank works (e.g. DeepSeek-V2-Lite has full-rank `q_proj` + low-rank KV): | Pattern | Projections | Already supported? | |---------|-------------|--------------------| | Full-rank Q | `q_proj` | Yes (`colwise_sharding_config`) | | Low-rank Q (MLA) | `q_a_proj`, `q_b_proj` | Yes (`_replicate_sharding_config`/`colwise_sharding_config`) | | Standard KV (GQA) | `k_proj`, `v_proj` | Yes (`colwise_sharding_config`) | | Low-rank KV (MLA) | `kv_a_proj_with_mqa`, `kv_b_proj` | Yes (`_replicate_sharding_config`/`colwise_sharding_config`) | | Output | `o_proj` or `dense` | Yes (`rowwise_sharding_config`) | | Q/K norms | `q_norm`, `k_norm` | Yes (`_replicate_sharding_config`, each independent) | | V-norm | `v_norm` | Yes (`_replicate_sharding_config`) | | DSA indexer | `indexer` | FSDP/EP only — **NOT under TP** (fails loud; the indexer's scatter_/index ops need the `spmd_types` backend, see `hf_sharding.py`) | The `_assert_all_states_sharded` backstop will fail loud at setup if any param/buffer-bearing attention (or other dense-path) module is left without a config, so an unhandled projection surfaces immediately. If found, report: "attention has unsupported projection `<name>`, needs `_set_layer_sharding_configs` update in `hf_sharding.py`." Also check: does the layer have `post_attention_layernorm`? (Some models omit it — already handled.) Does the embedding have extra buffers (e.g. Gemma4's `embed_scale`)? (Already handled — buffers are enumerated dynamically.) ### 1f. State dict adapter round-trip ```python from torchtitan.experiments.transformers_modeling_backend.state_dict_adapter import ( hf_to_titan_moe_state_dict, titan_to_hf_moe_state_dict, ) hf_sd = model.state_dict() titan_sd = hf_to_titan_moe_state_dict(hf_sd) roundtrip_sd = titan_to_hf_moe_state_dict(titan_sd) ``` Verify: all original keys present, no extra keys, values match. If keys are missing, the adapter needs new patterns for this model's key naming convention. ### 1g. Report Print a compatibility summary table: ``` Compatibility report for <model_id> model_type: <type> hidden_size: <H> moe_intermediate: num_experts: <E> top_k: <K> score_func: softmax|sigmoid route_norm: True|False shared_experts: none|additive|sigmoid-gated attention: gqa|mla|mla+dsa MoE layers: <indices> min_layers_needed: <N> state_dict_roundtrip: PASS|FAIL issues: none | <list> ``` ### 1h. Fix issues If issues were found, fix them in the experiment code now. Consult the "Issue Resolution Reference" at the bottom of this document for where each fix goes. The only true blocker is `nn.ModuleList` experts — for everything else, implement the fix and re-run the probe to verify. After all issues are resolved and the probe passes cleanly, proceed to Phase 2. ## Phase 2: Add Synthetic Test Config Add an entry to `_MODEL_CONFIGS` in `torchtitan/experiments/transformers_modeling_backend/.claude/skills/add_moe_model/scripts/numerical_equivalence.py`. ### Config rules Use production-scale `hidden_size` and `moe_intermediate_size` from the real config. Override everything else for a small single-GPU test: ```python "<model_type>": dict( model_type="<model_type>", hidden_size=<real>, # from pretrained config intermediate_size=<real>, # from pretrained config moe_intermediate_size=<real>, # from pretrained config num_local_experts=8, # reduced from real num_experts_per_tok=2, # reduced from real num_hidden_layers=<min>, # minimum for MoE layer to exist num_attention_heads=<fits>, # must divide hidden_size evenly num_key_value_heads=<fits>, # must divide num_attention_heads evenly vocab_size=256, max_position_embeddings=64, attn_implementation="sdpa", use_cache=False, # Include model-specific fields: # first_k_dense_replace, decoder_sparse_step, mlp_layer_types, # n_group, topk_group, n_shared_experts, norm_topk_prob, # routed_scaling_factor, q_lora_rank, kv_lora_rank, etc. ) ``` **Attention dimension rules:** - `num_attention_heads` must divide `hidden_size` evenly - `num_key_value_heads` must divide `num_attention_heads` evenly - For MLA models: use small `q_lora_rank` (16), `kv_lora_rank` (16), `qk_rope_head_dim` (8), `qk_nope_head_dim` (8), `v_head_dim` (16) - For DSA models: use small `index_n_heads` (2), `index_head_dim` (8), `index_topk` (8) **MoE layer selection rules:** - If `first_k_dense_replace > 0`: set it to 0 (force first layer to be MoE) - If model uses `decoder_sparse_step`: set to 1 - If model uses `mlp_layer_types`: set to `["sparse"]` **Expert group routing:** - If `n_group` is set, `num_local_experts` must be divisible by `n_group` and `num_local_experts / n_group >= 2` ## Phase 3: Run Numerical Equivalence **HF is the gold standard.** The numerical equivalence test compares Titan's MoE output against HF's output. The test must run HF's forward pass unmodified — do NOT patch, monkey-patch, or alter HF's computation to make the test pass. The only permitted modifications to close accuracy gaps are **temporary changes to titan core code**, which must be reverted after verification. Specifically: - Do NOT patch the HF model's activation function to match titan's - Do NOT replace HF's routing with titan's routing - Do NOT modify the test's comparison logic or thresholds - Do NOT skip or disable parts of the comparison - Do NOT cast HF tensors to different dtypes to hide precision diffs (the existing float32 gate cast is an exception — it aligns titan's autocast behavior, not the other way around) - If titan produces different results from HF, the gap analysis (Phase 3c) must identify WHY by temporarily changing titan to match HF, not by changing HF to match titan **Do NOT modify the pass/warn/fail thresholds.** The thresholds (KL < 1e-6 PASS, KL < 1e-3 WARN, KL >= 1e-3 FAIL) are calibrated for the known dispatcher precision gap. If a new model hits WARN, verify via Phase 3c that the diff is from known causes. If it hits FAIL, fix the integration, do not loosen the thresholds. ### 3a. Synthetic test ```bash CUDA_VISIBLE_DEVICES=<gpu> python \ torchtitan/experiments/transformers_modeling_backend/.claude/skills/add_moe_model/scripts/numerical_equivalence.py \ --models <model_type> ``` **Expected:** PASS with KL < 1e-6 and round-trip max_diff = 0.00. If it fails, diagnose: - `ValueError` from `_probe_hf_moe_block` — missing config attribute, fix the synthetic config - Shape mismatch in `load_state_dict` — state dict adapter needs update - `RuntimeError` from `grouped_mm` — check expert weight shapes - Numerical FAIL — investigate routing or weight transfer issue ### 3b. Pretrained test (if weights available) ```bash CUDA_VISIBLE_DEVICES=<gpu> python \ torchtitan/experiments/transformers_modeling_backend/.claude/skills/add_moe_model/scripts/numerical_equivalence_pretrained.py \ --model_dir <path> ``` **Expected:** PASS with KL < 1e-3 (known dispatcher precision gap). ### 3c. Verify numerical gaps via temporary core alignment (MANDATORY) **This step is NOT optional.** Execute it fully and automatically whenever any test shows WARN or higher. Do NOT ask the user before proceeding — apply patches, run tests, record results, report, revert. The entire sequence runs without user interaction. The goal: bring max_diff to 0.00 by temporarily making titan's implementation match HF's exactly. Every remaining diff must be experimentally explained — never claim a root cause without running code that proves it. **Rules:** - **Never give up on a patch because it crashes.** If a patch causes an error (e.g. dtype mismatch in scatter), fix the error by extending the patch. The patch is temporary — make it work, don't declare it "not applicable." - **Never claim a root cause without verification.** If you think "the diff is from bmm vs grouped_mm", prove it: make both sides use the same kernel and show the diff drops. If you can't change one side (e.g. HF rejects grouped_mm), change the other side (e.g. make titan use bmm temporarily, or manually compute the expert forward with the same kernel as HF using titan's weights). - **Adapt patches to the model's code path.** The known patches target specific code paths (e.g. f32 accumulation in `LocalTokenDispatcher.combine()`). If the model uses a different path, write an equivalent patch for that path. The principle is the same -- align titan's operations with HF's. - **Keep going until max_diff=0.00.** If the known patches don't get there, trace the forward step by step to find additional differences. **Procedure:** 1. **Record baseline** — note current max_diff and KL before any patches. 2. **Apply patches incrementally** — apply each applicable patch one at a time. After each, re-run the test and record the new max_diff. If a patch crashes, fix the crash (extend the patch as needed) and re-run. If a patch has no effect (e.g. top_k=1 makes reshape+sum trivial), record "no effect (top_k=1)" and move on. 3. **Handle kernel mismatches** — if the model doesn't support `grouped_mm` on the HF side, the test can't force the same kernel. In that case, verify experimentally: - Extract titan's expert weights after weight transfer - Run the expert forward manually using HF's kernel (e.g. `torch.bmm`) with titan's weights and the same input - Compare that output against HF's expert output - If they match (0.00), the remaining diff is confirmed as kernel difference, not a bug - Record this in the report as a verified kernel difference 4. **Investigate remaining diffs** — if max_diff is still non-zero after all patches and kernel alignment, trace the forward path step by step: router logits → routing weights → per-expert gate/up/down → score application → combine. Compare intermediate values between HF and titan to find where they first diverge. Make additional minimal patches to test hypotheses — if a patch doesn't help, revert it and keep digging. 5. **Report** — present a table to the user showing each patch and its measured impact: ``` Numerical gap analysis for <model_name> | Patch | max_diff | KL | Description | |-------|----------|----|-------------| | baseline | X.XXe-XX | X.XXe-XX | no patches | | +patch 1 | X.XXe-XX | X.XXe-XX | f32 accumulation in combine | | +patch 2 | X.XXe-XX | X.XXe-XX | reshape+sum (no effect, top_k=1) | | +kernel align | X.XXe-XX | X.XXe-XX | titan experts via bmm | | final | 0.00e+00 | X.XXe-XX | all diffs explained | ``` 6. **Revert ALL patches** — restore all core files: ```bash git checkout torchtitan/models/common/moe/token_dispatcher.py torchtitan/models/common/moe/moe.py ``` Verify with `git diff` that the revert is clean. The temporary patches are NOT part of the commit. **Known sources of numerical difference:** Do NOT copy-paste patches from this document — the titan core code may have changed since these were written. Instead, read the current code in `torchtitan/models/common/moe/token_dispatcher.py` and `torchtitan/models/common/moe/moe.py`, understand the principle behind each fix, and write a patch that applies cleanly to the current code. The three known differences and how to fix them: **1. bf16 accumulation before scatter_add** — `LocalTokenDispatcher.combine()` multiplies expert output by routing scores in f32 but casts back to bf16 before `deterministic_scatter_add`, so the accumulation happens in bf16. **Fix principle:** remove the cast to bf16, keep everything in f32 through the scatter_add, allocate an f32 output buffer, cast to the input dtype only at the very end. Example (illustrative, may not match current code): ```python # principle: remove .to(routed_output.dtype), use f32 buffer routed_output = routed_output.to(torch.float32) * scores out = scatter_add(f32_buffer, ..., routed_output) return out.to(x.dtype) ``` **2. scatter_add vs reshape+sum accumulation order** — `scatter_add` accumulates top_k contributions in expert-sorted order. HF uses `reshape(N, K, D).sum(dim=1)` which accumulates in token order (consecutive top_k entries per token). Different f32 summation order produces different results. **Fix principle:** store the argsort permutation from dispatch, in combine unsort back to original token order, then reshape and sum over the top_k dimension. For `top_k=1` this has no effect (only one element to sum). **3. topk sorted flag** — titan uses `torch.topk(sorted=False)`, HF defaults to `sorted=True`. With `sorted=True` topk values come back in descending order; the normalization sum is computed over differently ordered values, producing a small f32 rounding difference. **Fix principle:** change `sorted=False` to `sorted=True` in the router's topk call. For `top_k=1` this has no effect (only one element). With all three fixed, **all existing models produce max_diff=0.00.** If the new model still shows non-zero diff after all applicable fixes, the remaining difference is model-specific. Investigate as described in step 4 above. ## Phase 4: Parallelism Integration Tests Run the new model through the same parallelism configurations used for existing models. The integration test infrastructure is in `torchtitan/experiments/transformers_modeling_backend/tests/integration_tests.py`. ### 4a. Determine available GPUs ```bash nvidia-smi -L | wc -l ``` ### 4b. Run parallelism configs Test the new model with at least these configurations (adapt GPU counts to what is available). Add three temporary Python recipes to `config_registry.py`; configuration flags are not accepted by the launcher: ```python def _new_model_test_config(*, tp: int = 1, ep: int = 1) -> Trainer.Config: config = transformers_modeling_backend_debugmodel_moe() model_config = model_registry( "debugmodel_moe", seq_len=config.training.max_context_length, hf_model="<model_id>", ) config.model = model_config config.loss = CrossEntropyLoss.Config(global_vocab_size=model_config.vocab_size) config.training.steps = 2 config.parallelism.data_parallel_shard_degree = -1 config.parallelism.tensor_parallel_degree = tp config.parallelism.expert_parallel_degree = ep return config def new_model_fsdp() -> Trainer.Config: return _new_model_test_config() def new_model_fsdp_ep() -> Trainer.Config: return _new_model_test_config(ep=<EP>) def new_model_fsdp_tp_ep() -> Trainer.Config: return _new_model_test_config(tp=2, ep=<EP>) ``` Set `config.hf_assets_path` in the helper as well if the test needs the new model's tokenizer assets. Keep `training.steps = 2`: the goal is to verify the model initializes, parallelizes, and runs forward/backward without errors, not to converge. ```bash # FSDP only (minimum viable test) NGPU=<NGPU> MODULE=transformers_modeling_backend CONFIG=new_model_fsdp \ ./run_train.sh # FSDP + EP NGPU=<NGPU> MODULE=transformers_modeling_backend CONFIG=new_model_fsdp_ep \ ./run_train.sh # FSDP + TP + EP (if NGPU >= 4) NGPU=<NGPU> MODULE=transformers_modeling_backend CONFIG=new_model_fsdp_tp_ep \ ./run_train.sh ``` **Every config must be run.** There are only three — run all of them. Do NOT skip, do NOT mark as "NOT TESTED", do NOT mark as "SKIPPED". Every config in the final report must show either PASS or FAIL with a diagnosed error message and root cause. If a config fails: 1. Read the error traceback 2. Diagnose the root cause 3. If fixable in experiment code, fix it and re-run 4. If it's a fundamental limitation (e.g. DTensor mixing), report FAIL with the specific error and why it can't be fixed in experiment code **Redundant compute is a blocker.** If the only way to make TP work for a module is to run the full computation redundantly on every TP rank (e.g. all-gather input, run full forward, slice output), that is NOT an acceptable fix — it defeats the purpose of TP. If a model component (e.g. a custom attention mechanism) cannot be properly sharded, report it as a FAIL with the reason "no proper TP sharding available, would require redundant compute". Do NOT implement redundant-compute workarounds with pre/post hooks that all-gather and re-shard. **Use ShardingConfig, not hooks.** The HF backend sets ``_sharding_config`` on every sub-module so ``model.parallelize()`` handles all DTensor distribution and forward wrapping. Available helpers from ``torchtitan/models/common/decoder_sharding.py`` and ``torchtitan/experiments/transformers_modeling_backend/hf_sharding.py``: - ``colwise_sharding_config()`` — for linear projections (weight Shard(0), output Shard(-1)) - ``rowwise_sharding_config(output_sp=True)`` — for output projections (weight Shard(1), output Shard(1) for SP) - ``_hf_norm_sharding_config(enable_sp=True)`` — for norms (weight Replicate, activations Shard(1) when SP enabled) - ``_replicate_sharding_config(module)`` — for modules that should have Replicate params/buffers without TP sharding (dynamically enumerates params and buffers to avoid ``_shard_states`` errors) - ``ShardingConfig(in_src_shardings=..., in_dst_shardings=...)`` — for input redistribution at module boundaries (e.g. attention, MLP gather from Shard(1) to Replicate) Do NOT use ``register_forward_pre_hook`` / ``register_forward_hook`` to manually convert DTensor ↔ local tensors. Do NOT use ``parallelize_module`` from ``torch.distributed.tensor.parallel``. Use ``_sharding_config`` exclusively. Pass `<model_id>` to `model_registry` in the temporary recipes. If the base config has incompatible defaults for the new model (e.g. wrong tokenizer path), update those config fields in the helper as well. Do NOT give up because "the config is incompatible" -- fix the recipe. Common failure modes and fixes: - TP failures from unsupported attention projections — fix in `hf_sharding.py` by adding ShardingConfig in `_set_layer_sharding_configs` - EP failures from expert weight sharding — check FSDP `shard_placement_fn` - Shape mismatches from model-specific layer structure — fix in experiment code - Config incompatibility -- update the temporary recipe - Tokenizer mismatch -- set `config.hf_assets_path` to the model assets - Custom attention with no standard Q/K/V/O projections — FAIL, not fixable without redundant compute - Mixed DTensor / plain Tensor — a module has buffers not declared in its ShardingConfig's ``state_shardings``. Fix by using ``_replicate_sharding_config(module)`` which dynamically enumerates all params and buffers **Note:** If no GPUs are available, skip this phase and note it in the report. But if GPUs are available, all configs must be run and reported. ## Phase 5: Update Model Compatibility Doc Update `torchtitan/experiments/transformers_modeling_backend/MODEL_COMPATIBILITY.md` with the new model: 1. Add a row to the **Numerical Equivalence Summary** table with the model's status, KL, max_diff, cos_sim, and round-trip result. 2. Add a section under **Architectural Differences per Model** describing: - Router type and features (softmax/sigmoid, normalization, bias, scaling) - Expert format (standard/transposed, activation function) - Shared experts (none/additive/sigmoid-gated/layer-level) - Attention type (GQA/MLA/DSA/custom) - Any unsupported features and their impact - Summary of differences from Titan's MoE 3. Add a row to the **Parallelism Support** table (FSDP and FSDP+EP). 4. If the model has unsupported features requiring core changes, add or update entries under **Core Changes Needed for Full Support**. ## Phase 6: Lint and Commit 1. Run pre-commit on all changed files: ```bash SKIP=pyrefly-check pre-commit run --files <changed_files> ``` Fix any failures and re-run until clean. 2. Verify no core files were modified: ```bash git diff --stat # only files in experiments/transformers_modeling_backend/ ``` 3. Create commit: ``` Add <ModelName> to transformers MoE backend Adds synthetic test config with production-scale dimensions (H=<hidden_size>, I=<moe_intermediate_size>, E=8). Verified via numerical equivalence test (KL=<value>, max_diff=<value>). ``` ## Issue Resolution Reference Most issues are fixable in the experiment folder. Only `nn.ModuleList` experts is a true blocker requiring upstream changes. | Symptom | Where to fix | What to do | |---------|-------------|------------| | Expert weights transposed `(E, H, 2*I)` | `state_dict_adapter.py` | Add transpose in conversion; fix shape detection in `moe_replacement.py` | | `experts` is `nn.ModuleList` | **Cannot fix** — needs HF upstream | Report to user, stop | | Sigmoid scoring not auto-detected | `moe_replacement.py` | Add detection path in `_resolve_score_func()` (e.g. check router source for `sigmoid`) | | Router key `router.weight` not converted | `state_dict_adapter.py` | Add regex pattern in `_build_hf_to_titan_patterns()` | | Shared expert round-trip broken | `state_dict_adapter.py` | Fix reverse patterns in `_build_titan_to_hf_patterns()` | | MoE layer attr is `feed_forward` not `mlp` | `model.py`, `moe_replacement.py` | Handle alternative attribute name | | Attention projection not in ShardingConfig | `hf_sharding.py` | Add case to `_set_layer_sharding_configs` using `colwise_sharding_config()`/`_replicate_sharding_config()` | | Model rejects `grouped_mm` | `scripts/numerical_equivalence.py` | Skip `_experts_implementation` for this model | | `trust_remote_code` import error | Config loading | Try without `trust_remote_code` first | | `num_hidden_layers=1` has no MoE layer | Test config | Set `first_k_dense_replace=0` or increase layers | | State dict keys missing after round-trip | `state_dict_adapter.py` | Add patterns in `_build_hf_to_titan_patterns()` / `_build_titan_to_hf_patterns()` | | Layer-level MoE (router/experts are siblings of dense MLP) | `moe_replacement.py` | Treat dense MLP as shared expert inside the Titan MoE; replace `layer.mlp` with full MoE block containing router + experts + shared_experts(=dense MLP) | | TP fails with mixed DTensor/Tensor | `hf_sharding.py` | Module has undeclared buffers — use `_replicate_sharding_config(module)` which dynamically enumerates all params and buffers | | Embedding has extra buffers (e.g. `embed_scale`) | `hf_sharding.py` | Already handled — `set_hf_sharding_configs` enumerates embedding buffers dynamically | ## Files touched (experiment folder only) | File | Change | |------|--------| | `scripts/numerical_equivalence.py` | Add entry to `_MODEL_CONFIGS` | | `state_dict_adapter.py` | Add key patterns (only if round-trip fails) | | `hf_sharding.py` | Add ShardingConfig entries (only if new projection/norm names) | | `moe_replacement.py` | Add probing logic (only if new MoE pattern) |  --- name: inference_perf_hillclimb description: Profiler-driven hill-climbing to close the inference throughput gap between TorchTitan's unified model (running inside vLLM) and vLLM's native model. Benchmark with generate.py --benchmark, climb optimization rungs (compile / cudagraph / fused kernels), profile torchtitan vs the native target, then patch the single biggest gap at a time and re-measure. Use when the user wants to benchmark or optimize RL inference generation speed, reproduce previous hill climbing study, or invokes /inference_perf_hillclimb. --- # Inference Performance Hill-Climbing Close the gap between the TorchTitan unified model (in vLLM) and vLLM's native model on inference throughput, one profiler-found optimization at a time. Harness: `torchtitan/rl/generate.py --benchmark`. Append-only running log + full per-rung numbers live in `torchtitan/rl/docs/inference_gap_ablation.md`. Goal: match vLLM native 100%. ## Scope: close the IMPLEMENTATION gap, not the serving knobs This skill makes torchtitan's model run as fast as vLLM's NATIVE model at a FIXED model, workload, precision, and topology -- i.e. it closes the framework/overhead gap (target = native at the *same* config; ratio -> 1.0). It is explicitly NOT about the orthogonal levers that change the setup itself: - increasing batch size / seq_len / amount of input data - changing the model - lower precision (fp8 / quantization) - speculative decoding - disaggregated prefill - changing topology / number of GPUs Those move ABSOLUTE throughput but apply equally to native -- they don't change the torchtitan-vs-native ratio this skill chases. Hold them FIXED and IDENTICAL between torchtitan and native when measuring; otherwise you're no longer measuring the implementation gap. ## Before any run: confirm with the user (don't assume defaults) - **Model + parallelism**: which model, TP degree, SP/DP/PP/EP (e.g. Qwen3-32B, TP=8, SP off). Sets the sharding, all-reduce sizes, cudagraph capture sizes, and what "native" even is -- the dominant gap and right knobs change with it. - **Versions**: vLLM / torchtitan / PyTorch-nightly. They move fast and change kernels, cudagraph behavior, custom AR, and APIs. Record all three -- numbers across versions are not comparable. - **Workload**: input_len, batch_size, gen_tokens. Standard: W1 bs8/in1024/out128 (priority); W2 bs32/in4096/out1024 (needs `--max-seq-len 8192`). Sets prefill-vs- decode weighting, AR/capture sizes, memory -- the dominant gap depends on it (e.g. prefill capture helps short-gen W1 more than long-gen W2). - **compile + cudagraph** (the two settings that dominate perf; use AskUserQuestion): - compile: `aot_eager` (validated torchtitan default, per-layer torch.compile) | `inductor` (heavier) | `off`. **The fair target is native(eager)**, NOT native(inductor): inductor fuses the TP all-reduce into gemm/RMSNorm (`fuse_allreduce_rms`), eager does not -- it's ~+5% and not in the aot_eager path, so comparing to it invents a fake gap. - cudagraph mode: pick ONE for baseline AND target; never mix across rungs (see "cudagraph modes" under Harness). ## The hill-climbing loop (the core method -- PROFILE-DRIVEN, measure don't guess) **Every rung is justified by a PROFILE, never intuition**, and KEPT only if a re-benchmark moves throughput. The flag list is the residue of past profiles; the method is: profile -> single biggest gap -> patch exactly that -> re-measure -> re-profile. Two rules: (a) **trace time misleads** -- a huge all-reduce duration is usually cross-rank SPIN (sync wait), not reducible compute; confirm with tok/s. (b) **the biggest gap is often NOT a kernel** -- prefill cudagraph capture and arrival-spin were the top levers; profile broadly. 1. Confirm model / parallelism / versions / workload / compile+cudagraph (above). 2. Establish the ladder once, capture the **target**: `--native --compile aot_eager --cudagraph on` (same workload). 3. **PROFILE both** torchtitan-best and native (`--profile`). Decompose for the gap: - kernels: `/tmp/ablation_logs/analyze_trace.py` (single-rank diff: kernel us, extra kernels present in one but not the other, CPU-op deltas). - collectives: `/tmp/ablation_logs/analyze_ar_spread.py` over the 8 per-rank traces -- min(dur)=pure comm, spread(dur)=ARRIVAL SPIN. Single-rank kernel sums mislead for collectives. - host-bound: per-rank GPU busy% + idle-gap location (eager prefill vs captured decode). Low busy% on the driver rank => host-launch-bound. 4. Pick the SINGLE biggest gap (kernel / spin / host lag / capture coverage). 5. Implement exactly that as a new knob (custom_op patch in `kernel_ablation.py`, a `--model-2d` variant, or a cudagraph/config flag); often "make torchtitan do what native does here". 6. **Validate numerics**: example mode (no --benchmark) at TP>=2, greedy text bitwise-identical with vs without the patch. 7. Commit the rung (one git commit per table line). 8. Re-benchmark (W1, then W2); append to the doc. KEEP only if throughput moved (trace promise != throughput). 9. Repeat from step 3 (RE-PROFILE -- the biggest gap shifts after each patch) until within ~3% of target, or the remainder is structural (document it). ## Gap buckets: how to spot, what we found, the lever - **CPU / host overhead** (DTensor dispatch, ATen dispatch, per-launch Python). In an EAGER region the host can't feed the GPU -> it starves. Spot: per-rank GPU busy% < ~90% with idle gaps before compute kernels. Levers: cudagraph capture takes the CPU out of the per-kernel path (capturing PREFILL too was the single biggest lever, ~+20%); OR remove DTensor from the forward (`--model-2d` pure-local); OR cut per-launch overhead. Hidden under a captured graph -- only bites in eager prefill. NOTE: the structural lever is removing DTensor, not "2D" -- compile+cudagraph kill DTensor's CPU cost but its boundaries (from_local / subclass dispatch / placement views) fragment the captured graph -> launch jitter. Tensor rank is a red herring (3D==2D); register_sharding REGRESSES (-26%, more DTensor dispatch). - **Single kernel slowness** (same shape, torchtitan > native). Spot: `analyze_trace.py` per-kernel us. Found: NCCL all-reduce 23us vs vLLM custom one-shot AR 6.3us for small decode messages (a real algorithm difference); faster RoPE/attn variants. Lever: `--allreduce-vllm`, `--rope-kernel helion`, `--attn-backend flash`. Caveat: GEMM and RMSNorm are already at native parity (torch/Quack `rms_norm` even BEATS vLLM's) -- don't chase them. - **Fused vs separate kernels** (native fuses what torchtitan runs as N kernels: fused add+RMSNorm, fused QKV / gate-up GEMM, SiluAndMul). Spot: "extra kernels" in the torchtitan trace + higher CPU-op count. Lever: `fuse_qkv=True` + `fused_swiglu`, `--fused-addnorm`, `--silu-vllm`. Caveat: mostly a launch-count / host win -> matters in eager regions, shrinks under cudagraph. - **Communication straggler** (one rank reaches the collective late, the other 7 SPIN). Spot: `analyze_ar_spread.py` (min(dur)=comm==native, spread=arrival spin). Find WHAT BOUNDS the straggler: *compute?* per-rank compute time is usually UNIFORM, so no; *host-launch?* the DRIVER rank (rank 0) runs scheduler/sample/ output on its Python thread and falls behind in the eager prefill -> arrives last (busy 54% vs 94%, spin ~600us) -- the usual cause; fix = capture prefill (spin ~600us -> ~3us, the +20% lever) or lighten the per-step host path. *hardware/topology?* if the SAME rank straggles in NATIVE too, it's a GPU/NVLink-position effect, not fixable in the model. Rule: the AR kernel is innocent -- fix uneven ARRIVAL, not the collective. Operational gotchas: check GPU contention first (`nvidia-smi`; a co-tenant job spikes variance to +/-150 vs clean +/-1-7 -- re-run clean); NO Python-side logging/prints inside a patched forward (forces a torch.compile graph break); launch long runs detached (`setsid nohup ... &`) and clean GPU stragglers between. ## Harness **Build/restore**: stock `generate.py` on main is an EXAMPLE (single prompt). Bring the benchmark harness over from branch `ablation-inference`: `git checkout ablation-inference -- torchtitan/rl/generate.py` + the patch modules (`models/{kernel_ablation,qwen3_vllm_2d,helion_rope,vllm_fused_ops}.py`), fix the import (`from torchtitan.rl.examples.alphabet_sort import config_registry`), and add `rl_grpo_qwen3_32b` (mirror `rl_grpo_qwen3_14b`, `model_registry("32B")`, TP=8). `--benchmark` builds the engine OUTSIDE the timed region (excludes startup/compile/capture), runs `--warmup-runs` then times `--num-runs`, reports tok/s = batch*gen/wall; feeds synthetic token-ids (skips the tokenizer); sets `enable_prefix_caching=False`. Build the cudagraph `CompilationConfig` EXPLICITLY (the stock helper only emits `cudagraph_mode="full"`). **Launch**: `torchrun --nproc_per_node=<TP> generate.py --benchmark ...`. Env: conda `titan-rl` (see `torchtitan/rl/README.md` to build it). **Knobs** are the residue of past profiles -- NOT a required checklist, and NOT guaranteed to exist on a fresh checkout (the harness is rebuilt each time; re-add what you need, add NEW knobs for new gaps): - **Kernel patches** (torchtitan model; stack as the profile dictates): `--rope-kernel {helion,vllm} --silu-vllm --rmsnorm-vllm --allreduce-vllm --fused-addnorm --attn-backend {custom,flash}` -- custom_op monkeypatches in `models/kernel_ablation.py` / `vllm_fused_ops.py` (survive compile+cudagraph). FusedQKV+gate_up: use config `fuse_qkv=True` + `overrides/fused_swiglu.py` `@override`, not the superseded `--merged-gemm`. - **Whole-model paths** (mutually exclusive with kernel patches): `--model-2d {local,localfused,spmd,local3d,dtensor}`. `localfused` = pure-local (no DTensor) + fused add+RMSNorm (best); `spmd` = localfused but the TP AR routes through `spmd_types.redistribute(P->R)`. spmd_types is a pre-run CHECK (validates SPMD sharding via typechecking) -- keep it OFF for perf runs. - **compile / cudagraph**: `--compile {off,aot_eager,inductor}` `--cudagraph {on,off}` `--cudagraph-mode {full_decode_only,full}`. Capture PREFILL with `--cudagraph-mode full --max-num-batched-tokens --max-capture-size `, where P >= the prefill CHUNK size (= max_num_batched_tokens), NOT input_len. Decode capture sizes default to powers of 2 up to max_num_seqs (= batch). - **misc**: `--nccl-algo`, `--native` (target), `--profile` (per-rank chrome trace), `--max-seq-len` (long prompts). ### cudagraph modes (pick ONE for baseline AND target; never mix across rungs) These change WHO owns torch.compile: - **full_decode_only / full**: `CompilationConfig(mode=NONE)` -- vLLM does NOT compile; torchtitan's per-layer `aot_eager` is the only compile. full_decode_only = decode captured, eager prefill; **full** = also captures prefill (needs `--max-capture-size >= chunk`). KEEPS torchtitan compile. ## Example ladder + results ONE example trajectory, NOT a recipe -- each step was the biggest profiler-found gap at that point: baseline (eager DTensor, SP off) -> +compile(aot_eager) -> +cudagraph -> +tree all-reduce -> +Helion RoPE -> +FusedQKV/gate_up -> +SiluAndMul -> +FA3 -> +fused add+RMSNorm -> +vLLM custom AR -> +vLLM RMSNorm -> pure-local. Big jumps: cudagraph (0.03 -> 0.48x), vLLM custom AR (biggest KERNEL lever), dropping DTensor (pure-local). SiluAndMul / vLLM-RMSNorm / vLLM-RoPE were neutral-to-negative (torchtitan already fast). The biggest lever overall -- PREFILL cudagraph capture (~+20%) -- isn't a kernel; it was found later by profiling the residual arrival-spin straggler. Ratios vs native eager (32B TP=8, full_decode_only; full per-rung numbers and the FULL+prefill ladder are in the doc): | path | W1 (bs8/in1024/gen128) | W2 (bs32/in4096/gen1024) | |---|---|---| | DTensor ceiling (vLLM-AR + vLLM-RMSNorm) | 0.742 | 0.759 | | spmd_types, NCCL AR (`TT_SPMD_NCCL_AR=1`) | 0.698 | 0.796 | | spmd_types, vLLM AR (set_dist shim) | 0.846 | 0.877 | | pure-local (localfused) | 0.891 | 0.926 | | native | 1.000 | 1.000 | Capturing PREFILL too (`--cudagraph-mode full` + the max-batched-tokens/capture knobs) lifts every row ~+0.07-0.18x: DTensor ceiling -> ~0.90x, pure-local -> ~0.975x, spmd+vLLM-AR -> ~0.93x, at BOTH W1 and W2; native barely moves (its prefill is already lean). </yakaai_skills></article></body></html>

AI Agent Protocol & Instructions

Explore Related Technologies & Topics