../.claude/CLAUDE.md # XTuner Programming Standards ## Python ### Development Standards 1. When adding methods or functions to a module or class, declare **public interfaces** as early as possible and place **private interfaces** toward the end. 2. Python **>= 3.10** is required. Use modern syntax accordingly (e.g., `X | Y` union types). 3. Use `{}` instead of `dict()` for creating dictionaries. Use `[]` instead of `list()` for creating lists. 4. Use **double quotes** (`"`) for strings instead of single quotes (`'`). 5. Use `TypedDict` with `{}` syntax when a parameter requires a typed dictionary. 6. Prefer `pathlib.Path` over `os.path` for filesystem operations. 7. All new code must include **type hints** for function signatures (parameters and return types). 8. Maximum line length is **119** characters (configured in `pyproject.toml`). ### Docstring Standards 1. Use **Google Style** docstrings. 2. For classes, include `__init__` documentation directly under the **class docstring**; do not repeat it under the `__init__` method. 3. Parameter types **must** be declared in docstrings — include the type hint in parentheses immediately following the parameter name. 4. The **return type** must be explicitly declared in the docstring. 5. Do **not** declare raised exceptions in the docstring. 6. Do **not** declare Attributes in the class docstring. 7. If the return value is `None`, omit the `Returns` section. 8. Only provide docstrings for **public methods**. Private methods do not require docstrings. ### Linting, Formatting & Type Checking - **Ruff** for linting and formatting, **mypy** for type checking — see `pyproject.toml` for configuration. - `mypy` is strict for `xtuner/v1/*` (`ignore_missing_imports = false`). - Pre-commit hooks apply **only** to `xtuner/v1/` and `autotest/` — run `pre-commit run --all-files` before submitting a PR. - Docstrings are wrapped at **119** characters by `docformatter`. ### Error Handling 1. Never use bare `except:` — always catch specific exception types. 2. Avoid silencing exceptions with `pass` unless there is a documented reason. ### Testing 1. Every bug fix PR **must** include a regression test that reproduces the original bug. 2. Test functions must be named `test_` and grouped in classes named `Test`. 3. Use `pytest` as the test framework. Prefer `pytest.raises` for exception assertions. 4. Tests should be self-contained — avoid reliance on external services or GPUs unless explicitly marked with `@pytest.mark.gpu`. --- ## Git ### Branch Naming - Bug fixes: `fix/` - Features: `feat/` - Refactors: `refactor/` ### Commit Message Standards 1. Follow the style of previous commit messages to maintain a **consistent** project history. 2. Use the format: `[Tag] Short description` (e.g., `[Fix] Fix OOM in MoE expert routing`). 3. Common tags: `[Fix]`, `[Feature]`, `[Refactor]`, `[CI]`, `[Docs]`, `[Test]`. 4. **Never** include metadata or watermarks unrelated to the code itself, such as "generated by AI" or "Co-Authored-By: Claude." --- ## Pull Request Standards ### Bug Fix PRs 1. **Title**: `[Fix] ` 2. **Body** must include: - **Root Cause**: A clear explanation of why the bug occurred. - **Fix**: What was changed and why this is the correct fix. - **Test Plan**: How the fix was verified (unit test, manual reproduction, etc.). 3. Keep the diff **minimal** — fix only the bug. Do not refactor surrounding code, add unrelated improvements, or change formatting of untouched lines. 4. If the fix touches performance-critical paths (data loading, forward/backward pass, communication), include **before/after** benchmarks or justify why benchmarks are unnecessary. ### General PR Rules 1. One logical change per PR. Do not mix bug fixes with features or refactors. 2. All CI checks must pass before requesting review. 3. Resolve all review comments before merging — do not dismiss without explanation. --- ## Code Review Standards ### What to Check 1. **Correctness**: Does the code do what it claims? Are edge cases handled? 2. **Concurrency & Distributed**: Check for race conditions, deadlocks, incorrect `dist.barrier()` placement, and mismatched collective operations across ranks. 3. **Memory**: Watch for GPU memory leaks — tensors not freed, caches growing unbounded, missing `torch.no_grad()` in inference paths. 4. **Numerical Stability**: Flag potential issues like `log(0)`, division by zero, `float32`/`float16` overflow in loss computation or gradient scaling. 5. **Performance**: Unnecessary CPU-GPU synchronization (`item()`, `.cpu()` in hot paths), redundant data copies, inefficient collective communication patterns. 6. **API Contracts**: Do public interface changes maintain backward compatibility? Are deprecation warnings added for breaking changes? 7. **Resource Cleanup**: Are file handles, NCCL communicators, and CUDA streams properly cleaned up? ### Area-Specific Review Rules If a PR changes code under `xtuner/v1/rl`, read `.claude/rules/rl_review.md` before commenting or posting the final review summary. ### Review Output Format When performing code reviews, use the following structure: ``` ## Summary <1-2 sentence overview of the change> ## Issues ### Critical - [file:line] Description of the issue and suggested fix ### Warning - [file:line] Description of the concern ### Nit - [file:line] Minor style or readability suggestion ## Verdict APPROVE / REQUEST_CHANGES / COMMENT ``` ### Review Principles 1. **Be specific** — always reference the exact file and line number. 2. **Explain why** — don't just say "this is wrong"; explain the failure mode. 3. **Suggest a fix** — provide a concrete code suggestion when possible. 4. **Distinguish severity** — separate critical bugs from style nits. 5. Do not flag issues in code that is **outside the scope** of the PR diff. ## Rules Please refer to the docs in `.claude/rules` for the development guidelines. --- paths: - "xtuner/v1/datasets/**/*" --- # XTuner v1 Datasets Architecture Rules ## Core Data Contract (DataItem Protocol) Every dataset layer must produce/consume `DataItem`: ```python { "input_ids": list[int], # tokenized token IDs "labels": list[int], # same length; first token masked to -100 "num_tokens": int, # len(input_ids); used for sampler and packing } ``` - `labels[0] = -100` always (pretrain: first token not predicted) - `HardPackDataset.__getitem__` returns `list[DataItem]` (one per source sample in pack) - All other datasets return a single `DataItem` ## Tokenization Functions **Standard pretrain**: `PretrainTokenizeFunction` - Accepts: `{"content": "text"}` or `{"messages": [{"role": "pretrain", "content": "text"}]}` - Hash: `"{tokenizer_hash}_{source_hash}_{add_bos_token}_{add_eos_token}"` **Long text pretrain**: `LongTextPretrainTokenizeFunction` - Two-phase: cache phase returns chunk metadata, runtime phase tokenizes specific char ranges - Parameters: `chunk_size`, `tokenizer_chunk_chars`, `overlap_chars`, `min_chunk_tokens` - BOS only at `char_start==0`, EOS only at document end - Hash includes all chunking parameters ## JsonlDataset Architecture - **Shared memory**: entire file loaded once at init - **Offset table**: `np.ndarray` for O(1) line access - **Chunked mode**: activated when `chunks.npy` exists in cache - Expands `sampled` to per-chunk entries - `__getitem__` calls `tokenize_fn(data, char_start=cs, char_end=ce)` - `max_length` filter skipped (would incorrectly discard valid long-text samples) - `__len__` returns chunk count (not line count) ## Cache Directory Layout ``` cache_dir/ .xpuyu-cache-meta.json {file_xxhash}/ offsets.npy {tokenize_fn_hash}/ num_tokens.npy # always present chunks.npy # only for LongText mode: (total_chunks, 4) # [line_idx, char_start, char_end, num_tokens] ``` Cache key: `(file xxhash, tokenizer hash, tokenize_fn source hash)` - Tag-based lookup via `cache_tag` on `DatasetConfig` - Source hash auto-invalidates on code change ## Packing Strategies **HardPackDataset**: - Static pack_infos computed at init (deterministic from seed) - Returns `list[DataItem]` per pack - `global_pack=False`: per-dataset packing (default) - `global_pack=True`: concat all datasets then pack **PresetPackDataset**: - User-provided pack config (NPY directory + `paths.json`) - Validates structure and token ranges via vectorized numpy - Strategies for length mismatch: error/padding (short), error/truncate (long) - No state (deterministic from config file) ## Samplers **LengthGroupedSampler**: - Megabatch logic: group similar lengths, sort DESC, shuffle chunks - State: `{epoch, step, world_size, num_samples, total_size}` - `step` = `consumed_samples % total_size` **PresetSampler**: - User-provided order via `sampler_config_path` (`.npy` path only, mmap read) - Round-down to `global_batch_size * world_size` multiple (truncate tail; error if length too short) - Per-rank slice: `effective[local_rank::world_size]` on the truncated mmap-backed view (`self.global_order`) - `pack_level="preset"` in `DataloaderConfig`; requires `pack_config_path` and `sampler_config_path` - State dict compatible with existing resume logic ## Resume Contract ```python # Save state = dataloader.get_state_dict(consumed_samples) # → {"sampler": {...}, "dataset": {}} # Restore dataloader.load_state_dict(state) ``` - `consumed_samples` = total globally-seen samples (across all ranks) - Sampler converts to per-epoch `step` by modulo - Packing datasets have no state (deterministic) ## Multi-Dataset Support - `DatasetConfigList` = list of `{dataset, tokenize_fn}` pairs - `build_datasets()` returns `list[JsonlDataset]` (1 file → 1 dataset) - `sample_ratio` applied per-DatasetConfig via `sampled[]` repetition ## Key Implementation Files | File | Purpose | |------|---------| | `config.py` | Pipeline orchestration via `DataloaderConfig.build()` | | `pt_tokenize_fn/text.py` | `PretrainTokenizeFunction` | | `pt_tokenize_fn/long_text.py` | `LongTextPretrainTokenizeFunction` | | `jsonl.py` | `JsonlDataset` with shared memory + offset table | | `packing.py` | `HardPackDataset`, `ExpandSoftPackDataset` | | `sampler.py` | `LengthGroupedSampler`, `ParallelSampler` | | `preset_pack.py` | `PresetPackDataset` | | `preset_sampler.py` | `PresetSampler` | | `dataloader.py` | `Dataloader` wrapper | | `data_item.py` | `DataItem` schema | ## Design Principles 1. **Protocol-based**: All layers communicate via `DataItem` contract 2. **Cache-aware**: Tokenization cached with automatic invalidation 3. **Deterministic**: Packing and sampling reproducible from seed/config 4. **Resumable**: State dict pattern for checkpoint/resume 5. **Composable**: Tokenize → Dataset → Pack → Sample → Collate pipeline # RL Review Guidelines Use these rules whenever a PR changes code under `xtuner/v1/rl`. They extend the general review standards in `.claude/CLAUDE.md`; they do not replace them. ## Required Review Output Every review for an RL change must state the ProduceBatchResult impact near the top-level summary: ```text ProduceBatchResult impact: ``` Also include routed-experts impact when the PR touches routed-experts logic, rollout response handling, `RolloutState.extra_fields`, object references, or memory ownership around rollout outputs: ```text RoutedExperts impact: ``` Also include Ray concurrency impact when the PR touches Ray actor methods, actor construction, method decorators or wrappers, concurrency-group definitions, or rollout control paths: ```text Ray concurrency impact: ``` Do not leave these impacts implicit in the finding text. If a finding is about rollout status, pause/abort/timeout behavior, response handling, producer aggregation, routed-experts ownership, or Ray actor concurrency, repeat the relevant impact line inside that finding so the downstream effect is visible at the point of review. ## ProduceBatchResult Checklist Review whether the PR can change trainer-visible batch accounting, timing, or reward semantics in `ProduceBatchResult`. Check this area when the PR changes any of these paths: - `RolloutState.status`, `finish_reason`, or status conversion. - Abort, filter, expire, retry, timeout, cancellation, or failure handling. - Producer, agent loop, judger, replay buffer, rollout worker, rollout controller, or backend pause cleanup logic. - Writers, readers, tests, or fake implementations that construct or consume `ProduceBatchResult`, including trainer-facing code outside `xtuner/v1/rl`. Name the concrete field-level impact when any of these fields can change: - Batch status: `status`. - Returned groups: `rollout_states`. - Generation timing: `group_gen_count`, `group_gen_mean_s`, `group_gen_p50_s`, `group_gen_p99_s`, `group_gen_p99_p50_ratio`, `group_gen_pause_time_s`. - Replay-buffer leftovers: `leftover_init`, `leftover_completed`, `leftover_aborted`, `leftover_expired`, `leftover_failed`, `leftover_filtered`. - Reward accounting: `raw_rewards_sum`, `raw_rewards_count`. - Produced work counters: `produced_samples`, `produced_tokens`, `produce_time_s`. - Multi-task aggregation: `task_batch_sizes`, `task_results`. Common impacts to call out explicitly: - A sample moving between `ABORTED`, `FAILED`, `FILTERED`, `EXPIRED`, and `COMPLETED` can change the corresponding `leftover_*` counts. - Pause, abort, timeout, and cancellation changes can inflate or deflate `group_gen_*` timing, especially `group_gen_pause_time_s`. - Reward or filter-path changes can change `raw_rewards_sum`, `raw_rewards_count`, `produced_samples`, and `produced_tokens`. ## Ray Actor Concurrency Checklist Review whether high-volume rollout calls can starve pause, abort, health-check, restart, or shutdown RPCs. Check this area when the PR changes any of these paths: - Ray actor methods decorated with `@ray.method(...)` or `@ray_method(...)`. - Decorators or wrappers stacked on Ray actor methods, especially wrappers using `functools.wraps` or exposing `__wrapped__`. - Actor construction through `ray.remote(...)`, including `concurrency_groups` and `max_concurrency`. - Generate, pause, continue, abort, health-check, restart, or shutdown methods. When a Ray method is combined with another decorator, keep the Ray method decorator closest to the function definition: ```python @trace_rollout_endpoint("rollout.worker.generate") @ray.method(concurrency_group=ROLLOUT_CONCURRENCY_GROUP_GENERATE) async def generate(...): ... ``` Do not use the reverse order: ```python @ray.method(concurrency_group=ROLLOUT_CONCURRENCY_GROUP_GENERATE) @trace_rollout_endpoint("rollout.worker.generate") async def generate(...): ... ``` The current Ray actor-class processing unwraps decorated methods before reading Ray method metadata. If `@ray.method(...)` is outside a wrapper, metadata such as `__ray_concurrency_group__` can remain on the discarded wrapper, causing the method to fall back to the default concurrency group. Review the effective dispatch behavior, not only the visible decorator order: - After following `__wrapped__`, the callable inspected by Ray must retain the expected Ray method metadata. - Every named method concurrency group must be declared by the corresponding `ray.remote(...)` actor construction. - High-volume `generate` methods must not accidentally fall back to the default group. - Pause, abort, and other control RPCs must remain schedulable while generation is saturated. - Slow health or lifecycle methods must not occupy every slot needed by critical control RPCs. - Blocking work inside async actor methods must not block the actor event loop or concurrency lane. ## RoutedExperts Checklist Review whether routed-experts ownership and cleanup remain correct. Check this area when the PR changes any of these paths: - LMDeploy rollout response handling. - `return_routed_experts`, `routed_experts`, `RolloutState.extra_fields`, or object-ref plumbing. - Abort, cancellation, timeout, filter, retry, or failure paths before response handling completes. - Background tasks, replay-buffer storage, trainer batches, metrics, tests, or fake rollout responses that can retain routed-experts object refs. For LMDeploy rollout, `rollout_worker` obtains routed-experts object refs from the LMDeploy shared store. At that point, ownership moves from LMDeploy to XTuner. A review finding should call out both sides of the ownership boundary when relevant: - Leak before transfer: requests whose routed experts remain in LMDeploy because XTuner never obtains the object refs. - Leak after transfer: object refs that XTuner keeps alive too long through `RolloutState`, replay buffer, trainer batches, metrics, fake tests, or background tasks. --- name: add_hf_model description: > Add support for a HuggingFace model in XTuner so it can be trained with the full set of parallelism strategies and optimization switches. The user can provide a HuggingFace Hub repo id, a local model directory, or a model family already supported by `transformers`. This skill walks through the full path: locating the reference implementation, classifying the model, splitting new code across the model / module / ops layers, implementing the XTuner model class and config, registering the entry point, validating bitwise numerical parity with HuggingFace, writing regression tests, and enabling the training optimizations (EP/SP, micro-batch, torch.compile, fp8, activation offload). --- # Add a HuggingFace model to XTuner A complete handbook for porting any HuggingFace causal LM into XTuner with correct training behavior. Covers four buckets: **Dense LLM**, **MoE LLM**, **VLM / compose model**, and **`trust_remote_code` model**. Every step below is grounded in concrete `file:line` references in this repository — read the cited code before writing yours. > **Anti-spaghetti reminder** — XTuner separates *what the model is* (`*Config`), > *what it computes* (`Dense` / `MoE` / `BaseComposeModel` subclass), and *how its > weights map to HF safetensors* (`to_hf_key_list`, optional `safetensors_to_params` > / `param_to_safetensor`). Keep new code inside those three seams. **Never** add > `if model_name == "..."` branches to the base classes. --- ## 0. Locate the source code The only goal of this step is to find the **HuggingFace reference implementation you will port from**, and to note the `model_type` you will register under. Do not pick attention/router/layer here — that happens naturally in §1 as you read the code. 1. If the user gives a **remote-code** model, the modeling code ships in the repo as `modeling_*.py` — read it directly. 2. If the user gives **weights only**, confirm the target Python environment with the user, then find the matching implementation inside that env's `transformers` install. If no code exists for that `model_type`, stop and ask the user. 3. Confirm the weights are **training-suitable** — at least bf16 — so training precision can be aligned later. Flag fp8/int4-only checkpoints to the user. 4. Note two strings from the HF config — you will need both downstream: ```python from transformers import AutoConfig cfg = AutoConfig.from_pretrained(path, trust_remote_code=True) cfg.model_type # registration key (§4) cfg.architectures[0] # the modeling class you are mirroring ``` --- ## 1. How to do the migration XTuner splits a model into three layers. You port by reusing what already exists at each layer and only adding what is genuinely new — so the work is "read the HF modeling code, then match each of its pieces to the XTuner component below". The locations are listed plainly; you do not need extra guidance to map your model's attention / router / layer onto the matching one. - **model layer** (`xtuner/v1/model`) — the model body. - Dense body: `xtuner/v1/model/dense/dense.py` - MoE body: `xtuner/v1/model/moe/moe.py` - Multi-modal (compose) bodies: `xtuner/v1/model/compose/` — reuse the language model and add only the ViT / projector pieces. - **module layer** (`xtuner/v1/module`): - Decoder layers: `DenseDecoderLayer`, `MoEDecoderLayer` — `xtuner/v1/module/decoder_layer`. Small changes add a switch on the existing layer; a genuinely novel architecture may add `{model_name}_decoder_layer`. - Attention configs: `MHAConfig` (`xtuner/v1/module/attention/mha.py`), `MLAConfig` (`.../mla.py`, DeepSeek-style latent), `GatedDeltaNetConfig` (`.../gated_deltanet.py`, linear). - Router configs (MoE): `GreedyRouterConfig` (`xtuner/v1/module/router/greedy.py`, softmax top-k — Qwen3 MoE, GptOss), `NoAuxRouterConfig` (`xtuner/v1/module/router/noaux_router.py`, grouped sigmoid aux-loss-free — DeepSeek V3). - Also: rope, rms_norm, etc. - **ops layer** (`xtuner/v1/ops`) — kernels such as attention and rms_norm. ### Existing models to copy from Pick the one whose attention + (router) match yours; the closer it is, the smaller your diff. | Your model is… | Copy from | |--------------------------------------|--------------------------------------| | Dense LLM | `xtuner/v1/model/dense/qwen3.py` | | MoE LLM | `xtuner/v1/model/moe/qwen3.py` | | MoE with MLA + aux-loss-free router | `xtuner/v1/model/moe/deepseek_v3.py` | | MoE needing weight-layout transform | `xtuner/v1/model/moe/gpt_oss.py` | | Compose / VLM | `xtuner/v1/model/compose/qwen3_vl/` | After you have a plan, **write an integration design doc into `docs/design/model`** and confirm with the user that the decomposition is reasonable before implementing. ### File layout for the new model | Bucket | Model file | Public re-export | Test file | |----------------------|-------------------------------------------------------------|-----------------------------------------------|--------------------------------------------| | Dense LLM | `xtuner/v1/model/dense/.py` | `xtuner/v1/model/__init__.py` | `tests/model/test__dense.py` | | MoE LLM | `xtuner/v1/model/moe/.py` | `xtuner/v1/model/__init__.py` | `tests/model/test__moe.py` | | Compose / VLM | `xtuner/v1/model/compose//...` | `xtuner/v1/model/__init__.py` | `tests/model/test_.py` | --- ## 2. Implement the model class Subclass `Dense`, `MoE`, or `BaseComposeModel`. You typically only need to override one method. ### 2.1 `to_hf_key_list(self, key: str) -> list[str]` — **mandatory, abstract on `BaseModel`** Declared at `xtuner/v1/model/base.py`. Translates an **XTuner-side** parameter name into one or more **HF-side** safetensors keys. Return-list cardinality: - **1 → 1** for plain weights (most params). - **1 → N** for fused MoE experts (one fused param explodes into N per-expert HF keys; see `xtuner/v1/model/moe/qwen3.py`). Patterns you almost always need (study `xtuner/v1/model/dense/qwen3.py` and `xtuner/v1/model/moe/qwen3.py`): ```python # 1. Tied embeddings: redirect lm_head → embed_tokens before adding the model. prefix. if self.config.tie_word_embeddings and "lm_head" in key: key = key.replace("lm_head", "embed_tokens") # 2. Add the "model." prefix that HF wraps everything in. if "layers" in key or "embed_tokens" in key: key = "model." + key # 3. HF MoE nests experts/gate under .mlp.; XTuner does not. if "layers" in key: key = re.sub(r"layers\.(\d+)\.(experts|gate)", r"layers.\1.mlp.\2", key) # 4. Top-level norm. if key.startswith("norm."): return [key.replace("norm.", "model.norm.")] # 5. MoE expert fusion: one fused param → N HF keys. if "fused_w1w3.weight" in key: out = [] for i in range(self.config.n_routed_experts): out.append(key.replace("fused_w1w3.weight", f"{i}.gate_proj.weight")) out.append(key.replace("fused_w1w3.weight", f"{i}.up_proj.weight")) return out # 6. Model-specific buffers (e.g. FoPE: rotary_emb.sin_coef / cos_coef are persistent buffers). if key.startswith("rotary_emb."): return [key.replace("rotary_emb.", "model.rotary_emb.")] ``` A prefix-region remap that differs by *deployment* (e.g. a text tower nested under a VLM as `model.language_model.`) does **not** belong here — put it in the config's `hf_key_mapping`. See §5.1. ### 2.2 `safetensors_to_params` / `param_to_safetensor` — **optional, only for layout mismatches** If the HF storage layout cannot be expressed as “same tensor, different name”, you must transform the bytes. Canonical example: GptOss stores expert weights as `(num_experts, hidden_size, expert_dim * 2)` but XTuner fuses them as `(num_experts * 2 * expert_dim, hidden_size)` — see overrides at `xtuner/v1/model/moe/gpt_oss.py`. Use this hook only when transposing/reshaping is unavoidable. Renames go in `to_hf_key_list`; numerics go here. ### 2.3 Do **not** override - `build_layers`, `build_embeddings`, `_init_weights` — the base handles them via the config. Override only if the new model adds a layer type the base cannot express (rare). - `forward` — same reasoning. If you find yourself touching it, re-evaluate the bucket choice; the problem is probably in the config (`attention`, `router`, `first_k_dense_replace`). --- ## 3. Implement the config Two layers: a **base config** that reads HF and a **size-specific subclass** per released checkpoint. ### 3.1 Base config — three required members Mirror `Qwen3MoEConfig` (`xtuner/v1/model/moe/qwen3.py`) or `Qwen3DenseConfig` (`xtuner/v1/model/dense/qwen3.py`). 1. **`build(self) -> `** — return `(self)`. 2. **`@classmethod from_hf(cls, hf_path) -> Self`** — read the HF config and map its fields one-to-one onto your config. The field set is whatever the live config class declares (`TransformerConfig` in `xtuner/v1/model/base.py`, `MoEConfig` in `xtuner/v1/model/moe/moe.py`); read it, don't work from a hard-coded list. A few non-obvious traps: - Use `RopeParametersConfig.from_hf_config(hf_config)` from `xtuner/v1/module/rope/rope.py`. It already handles both the `rope_parameters` dict (HF ≥ 5.2.0) and the legacy `rope_scaling` dict (HF 4.57.x), plus YARN / FoPE / mrope special fields. **Do not** parse rope yourself. - For optional HF fields, use `getattr(hf_config, "", )` rather than `hf_config.` — older checkpoints drop fields. 3. **`@property hf_config(self) -> | None`** — the inverse of `from_hf`: every field `from_hf` reads must be re-emitted so `save_hf` round-trips. **Built-in vs. `trust_remote_code` — the only axis that changes `from_hf` / `hf_config`:** a `trust_remote_code` model is structurally still Dense / MoE / VLM; the *only* difference is whether `transformers` ships a built-in `Config` class. Decide which case you are in and follow the matching column — this is the single place that contract is defined. | | Built-in config (e.g. Qwen3) | `trust_remote_code` (no built-in config) | |---|---|---| | `from_hf` reads via | `.from_pretrained(hf_path)` | `AutoConfig.from_pretrained(hf_path, trust_remote_code=True)`, then `assert hf_config.model_type == ""` | | `hf_config` returns | a populated `` re-emitting every field | `None` | | `save_hf` weight-map | written from `hf_config` | same | | `save_hf` side files | derived from `hf_config` | falls back to copying `config.json` / tokenizer / `*.py` from `self._hf_path` so the dir stays loadable with `trust_remote_code=True`; module-cache test in the matrix below | Compare `Qwen3MoEConfig.hf_config` (built-in) and `Qwen3MoEFoPEConfig.hf_config` (returns `None`), both in `xtuner/v1/model/moe/qwen3.py`. **Do not** invent a built-in HF config class for a remote-code model — return `None`. **Invariants:** - `from_hf` ↔ `hf_config` must round-trip. The `test_save_hf` test (§7) is the enforcement mechanism — every key in the original index must appear in the saved index with byte-equal tensors. - The config must be **pickleable and side-effect free**. No env vars, no file I/O, no `torch.cuda.*` at construction time. ### 3.2 Size-specific subclasses Add one subclass per released checkpoint that hard-codes the published dimensions. Mirror `Qwen3MoE30BA3Config` / `Qwen3MoE235BA22Config` in `xtuner/v1/model/moe/qwen3.py`. Even if only one size ships, define at least one — it is the entry point for `get_model_config` and for tests. --- ## 4. Register the entry point Open `xtuner/v1/model/__init__.py` — all four edits are short: 1. **Import** — bring in the new config classes. 2. **Alias** — add `"": ()` to `model_mapping`. Lookup normalizes case and `-`/`_`, so `"my-model-7b"` and `"my_model_7b"` are the same key. 3. **Dispatch** — add `elif cfg.model_type == "":` to `get_model_config_from_hf` returning `.from_hf(model_path)`. The catch-all `raise ValueError(...)` must stay last. 4. **`__all__`** — re-export the new public config names. **Never** key dispatch on parameter names, file paths, or architecture strings. Only `model_type`. --- ## 5. Compose / VLM specifics Compose models live under `xtuner/v1/model/compose//` and follow a three-tower pattern. See `xtuner/v1/model/compose/base.py` and `xtuner/v1/model/compose/qwen3_vl/qwen3_vl_config.py`. ### 5.1 Config shape `BaseComposeConfig` (`compose/base.py`) carries three sub-configs: ```python vision_config: XTunerBaseModelConfig # e.g. Qwen3VLVisionConfig projector_config: XTunerBaseModelConfig # e.g. Qwen3VLProjectorConfig text_config: XTunerBaseModelConfig # any Dense/MoEConfig subclass freeze_vision: bool = False freeze_projector: bool = False freeze_language: bool = False ``` Reuse existing text-tower configs where possible. `Qwen3VLDense*` delegates to the Dense path via `xtuner/v1/model/dense/qwen3vl_text.py`. **Do not** copy-paste the entire dense/MoE config — compose over them. **Reused text-tower key prefixes — use `hf_key_mapping`, not `to_hf_key_list`.** A text tower reused inside a VLM keeps its *standalone* `to_hf_key_list` (it emits `model.<...>` as if it were a plain LLM). The compose checkpoint nests it one level deeper (`model.language_model.<...>`). Do **not** teach the text tower's `to_hf_key_list` about that prefix — that leaks the compose context into a tower that must also work standalone. Set the remap on the **text config** instead: `hf_key_mapping = {r"^model\.": "model.language_model."}` (applied in `_init_load_spec`). Keep *structural* renames in `to_hf_key_list`; the deployment-dependent prefix region belongs in `hf_key_mapping`. ### 5.2 Model shape `BaseComposeModel` (`compose/base.py`) builds three sub-models: `vision_tower`, `multi_modal_projector`, `language_model`. - **`from_hf(hf_path)`** — delegates to each sub-model with `strict=False` (vision keys are missing from the language-model state dict and vice versa) and unions the missing-keys sets. - **`save_hf(hf_dir)`** — saves each tower with a distinct prefix (`"model-language"`, `"model-vision"`, `"model-projector"`) and merges the three `weight_map`s into one `model.safetensors.index.json`. ### 5.3 `hf_config` for VLMs Many VLMs ship without a stable HF top-level config class; `hf_config` returns `None` (`compose/qwen3_vl/qwen3_vl_config.py`). In that case, `save_hf` falls back to copying source files from `self._hf_path` (see §3.1). ### 5.4 Vision-side quirks to watch for - `attn_impl` selectable per checkpoint: `"flash_attention" | "flex_attention" | "eager_attention"`. - `deepstack_visual_indexes` (e.g. `[8, 16, 24]`) — auxiliary supervision depths on the vision tower; must round-trip. - Specialized rope: `rope_type="qwen3_vl"`, `mrope_section=[24, 20, 20]` — delegated to `RopeParametersConfig.from_hf_config`. - The projector’s `torch.compile` path is enabled on torch ≥ 2.9.1 (see commit `f6d74efb`). Tests must validate both compiled and eager outputs match. --- ## 6. Baseline parity — no complex parallelism yet First make the model correct on the **simplest execution path**: single rank, no EP/SP, no `torch.compile`, no fp8, no activation offload. Align precision from the inside out — ops → decoder layer → full model — reusing XTuner's own modules at every level. Only once this baseline holds do you add the parallel and optimization features (§8). 1. **Reuse, don't reimplement.** Use the existing XTuner modules — MHA, MLA, GatedDeltaNet, GroupLinear, etc. Do not re-implement them with naive torch. 2. **Route op differences through `XTUNER_HF_IMPL`.** This env var already exists (`xtuner/v1/utils/misc.py`) and selects HF-exact ops via per-op selectors: `xtuner/v1/ops/attn_imp.py::get_attn_impl_fn` (eager attention), `xtuner/v1/ops/rms_norm/__init__.py::get_rms_norm_fn` (native torch rms_norm), and `xtuner/v1/ops/gated_deltanet/__init__.py::{get_chunk_gated_delta_rule_fn, get_causal_conv1d_fn}` (return the canonical fla / causal_conv1d wrappers HF calls instead of XTuner's compile-friendly `torch.library.custom_op` wraps — same forward output, but the backward op graph now matches HF). Selectors read the env var **live** so a test can toggle it per model instance. If your model introduces a new op whose fast path is not bitwise against HF (typical for fla-backed linear-attention variants), add its HF-exact branch in the corresponding op selector the same way. **You may only patch at the ops level** — never take the shortcut of patching the entire `XTunerDecoderLayer.forward`. Pair `XTUNER_HF_IMPL` with `XTUNER_DETERMINISTIC=true`. The latter gates an autotune-pin block in `xtuner/v1/__init__.py` that monkey-patches `triton.autotune` to lock every kernel to its first config and disable the result cache. Without the pin, fla's `@triton.autotune` picks tiling / reduction order per kernel and the choice drifts between runs, producing 1 ULP per linear-attention layer that the chain rule amplifies to ~1.0 at the model boundary (§12). `tests/conftest.py` sets the env var and runs `import xtuner.v1` so the patch installs before fla is imported. The patch lives on the `xtuner.v1.*` import chain (not in the conftest) because `multiprocessing.spawn` children of `MultiProcessTestCase` don't go through pytest's conftest — they re-import the test class top-level (which pulls `xtuner.v1.model.*`) and the patch must run in those processes too. Production training doesn't set the env var, so the block is inert there. 3. **Decoder-layer bitwise parity — forward AND backward, both required.** Build one XTuner `DecoderLayer` (one per distinct layer type — e.g. a linear and a full layer for a hybrid model) and the matching HF layer, set `XTUNER_HF_IMPL` and `XTUNER_DETERMINISTIC`, feed identical inputs, and align **bitwise** — not a tolerance. **Backward is not optional**: a `single layer → final norm → lm_head → CE loss → loss.backward()` test asserts the **layer output, the loss, and the input gradient `dL/dx`** all bitwise-equal to HF — test one full and one linear layer separately (reference: `test_decoder_layer_bitwise_parity`). A bitwise *forward* does not imply a bitwise backward — the forward values can round identically while the backward graph differs (see the fp32-softmax note in §12). This is the **primary** path, required at all scales, and the only feasible one for large models (it never loads the full model — see "load only what you use" below). When you hand-build a single layer's inputs, the two sides' `causal mask` / `seq_ctx` must match exactly — see §12. **Load only what you use** so the test fits any model size: build the XTuner tower on `meta` and materialize (`to_empty` + selective load) *only* the tested layer + `norm` + `lm_head` via `HFCheckpointLoader.load(key)` (it reads just the needed safetensors shard); build a *standalone* HF `DecoderLayer` + norm + lm_head and load the same keys. Never `from_pretrained` the whole model for a single-layer test. (`build_rotary_embedding` builds the rotary on CPU with real buffers even under meta, so `model.rotary_emb` is usable without materializing the layers.) For < 40B you may also read per-layer hidden states from one full-model forward (HF `output_hidden_states` + XTuner `return_hidden_states`); mind the off-by-one in §12. 4. **Model config matches HF hyperparameters** — see the §3.1 built-in vs. remote-code contract for how `from_hf` / `hf_config` differ between the two. 5. **Full-model forward parity for models < 40B**, bitwise against HF, on the single-rank path (FSDP on or off; no EP / compile / fp8 yet). 6. **`save_hf` / `from_hf` correctness** (round-trip; see §7). 7. **Loss-convergence trace + engine test.** Once model forward and decoder-layer forward/backward are aligned, record a loss-convergence trajectory and add an engine test case. **Commit the baseline — do not stop to ask.** Committing is part of this workflow: invoking this skill authorizes the in-workflow commits. Once parity holds on the simple path and lint/tests pass, commit immediately so this correct baseline is preserved before any parallel/optimization work begins (the §8 sub-agents branch off this commit). Commit each logical step as you finish it (§10) rather than batching everything to the end or waiting for a separate go-ahead. Report back to the user: whether bitwise parity was achieved, and the residual error (decoder-layer level and model level) once XTuner's internal components are used. --- ## 7. Baseline tests — add these first Before any parallel feature, add the most basic regression tests so the baseline (§6) is locked and the §8 work has a fixed reference to measure against. Add `tests/model/test__.py`, mirroring `tests/model/test_qwen3_moe.py` (MoE) or `tests/model/test_qwen3_dense.py` (Dense). All tests extend `xtuner._testing.DeterministicDDPTestCase` for deterministic distributed setup. Reference test cases: - decoder-layer bitwise parity + save_hf round-trip: `tests/model/test_qwen3_5_dense.py` (hybrid linear/full dense VLM; toggles `XTUNER_HF_IMPL` per case) - model forward & save & load parity: `tests/model/test_qwen3_moe.py` - engine training test: `tests/engine/test_moe_train_engine.py` ### 7.1 Baseline test matrix All single execution path — no EP / dispatcher / compile / fp8 here; those are §8 tests, written against this baseline. | Case | Dense | MoE | Compose | Notes | |-------------------------------------------------|:-----:|:---:|:-------:|------------------------------------------------------------------------------------| | Decoder-layer parity — output + loss + `dL/dx` | ✅ | ✅ | ✅ | **Required at all scales, bitwise.** Test one full and one linear layer *separately*. Standalone & memory-light: build the XTuner tower on `meta` and `materialize_xtuner_submodule` only {tested layer, `norm`, `lm_head`}; build a standalone HF `DecoderLayer`; load via `HFCheckpointLoader`. Under `XTUNER_HF_IMPL` assert the **layer output, loss, and `dL/dx` all bitwise** (forward AND backward at the layer level — a bitwise forward ≠ bitwise backward, see §12). Reference: `test_decoder_layer_bitwise_parity`. | | Whole-model forward + backward parity (by scale) | ✅(4B)| scale | ✅(4B) | Run the **top-level model** end-to-end vs HF, eager both sides (`XTUNER_HF_IMPL` + `XTUNER_DETERMINISTIC` + HF `eager`). For a **VLM** that's the compose model on an **image** prompt (vision + projector + text — subsumes "VLM forward parity"); for a plain LLM, the text/LM model. With both env switches on plus the autotune pin loaded (§6.2), **forward AND backward are bitwise**: assert logits + loss + `dL/d(pixel_values)` / `dL/d(inputs_embeds)` all `== 0.0`. **Scale**: small = fwd+bwd; medium = forward only; large (can't e2e-forward) = skip → integration. Reference: `test_vl_forward_parity`. | | Vision-tower bitwise parity | — | — | ✅ | **Compose/VLM only.** Compare HF's `visual` output (`pooler_output`, which is **post-merger**) bitwise (0.0), both eager. XTuner splits that boundary into `vision_tower` (patches → merged hidden) + `projector` (the merger MLP → text dim), so load **both** standalone to match it — don't build the whole VLM just for the vision path. Build them on the **real device, not `meta`**: the vision rotary is computed into a buffer at forward time, so `to_empty` would leave garbage (unlike the text rotary, which `build_rotary_embedding` builds real even under meta — §6.3). Vision attention is *non-causal* → `eager_attention(causal=False)`, which `XTUNER_HF_IMPL` selects. Reference: `test_vision_tower_bitwise_parity`. | | FSDP forward parity | ✅ | ✅ | ✅ | Looser tolerance (≈ 3e-2 is the established budget). | | `save_hf` round-trip (byte-equal tensors) | ✅ | ✅ | ✅ | See `tests/model/test_qwen3_moe.py`. | ### 7.2 Parity bar by model scale Two independent axes: **Per-layer (op) parity — required at _all_ scales, bitwise.** Test one full and one linear decoder layer separately; for each assert layer **output + loss + `dL/dx`** bitwise vs HF under `XTUNER_HF_IMPL`. Keep it memory-light (XTuner `meta`-build + materialize only the tested layer + `norm` + `lm_head`; standalone HF layer; load via `HFCheckpointLoader`) so it runs even for models too large to forward end-to-end. This is the bitwise guarantee for both forward and backward ops. **Whole-model (integration) parity — on the _top-level_ model, graded by what fits one GPU.** Run the real model end-to-end: for a VLM that's the **compose model on an image prompt** (one test covers vision + projector + text — don't add a separate text-tower model test); for a plain LLM, the text/LM model. - **Small** (e.g. dense 4B): both forward (logits/loss) AND backward (`dL/d(pixel_values)` for a VLM, `dL/d(inputs_embeds)` for text) **bitwise**. Prerequisites: `XTUNER_HF_IMPL=true` so every op selector returns the HF-canonical callable (§6.2) AND `XTUNER_DETERMINISTIC=true` so Triton autotune is pinned (§6.2). With both on, even hybrid linear/full models match HF byte-for-byte through the whole compose chain. Don't check weight gradients — they accumulate across positions + tied `lm_head` and aren't the right granularity for parity. - **Medium**: forward only (logits/loss bitwise). - **Large** (cannot forward end-to-end on the test GPUs): skip the whole-model test; leave it to integration tests. The per-layer test still gives op-level forward+backward bitwise. ### 7.3 Required test idioms `tests/model/test_qwen3_5_dense.py` is the reference template; extend `DeterministicDDPTestCase` (`xtuner._testing`), which provides the shared parity scaffolding so the test file only holds model-specific logic: - `with self.hf_impl():` — sets/restores `XTUNER_HF_IMPL=true`; build the XTuner model **inside** the block so its attention picks the eager op at construction. - `self.materialize_xtuner_submodule(model, submodule, prefix, loader)` — for a model built on `meta`, materialize *only* that submodule on CUDA and load its weights (so a single-layer test fits any model size). - `self.load_params_from_hf(module, key_for, loader)` / `self.xtuner_ckpt_key(model, name)` — selective per-param load by checkpoint key (`HFCheckpointLoader` reads only the needed shard); `xtuner_ckpt_key` applies `to_hf_key_list` + the config's `hf_key_mapping`. ```python # Construct on meta, materialize on device — single-rank baseline, eager, no EP. with torch.device("meta"): cfg = get_model_config_from_hf(hf_model_path) cfg.compile_cfg = False # baseline tests run eager model = cfg.build()._to_device_dtype(dtype=torch.bfloat16, skip_buffers_dtype=True) # Optionally shard with FSDP (data-parallel only), then load. model.fully_shard(fsdp_config=FSDPConfig(cpu_offload=False)) model.from_hf(hf_model_path) ``` Older `patch_hf_rms_norm` / `patch_hf_rope` helpers also live in `xtuner._testing`. **Check they actually fit your model before using them** — they were written for earlier models and match HF modules by class-name substring and attribute name (e.g. `variance_epsilon`). On a newer model they can silently mis-patch (wrong eps attribute, or replacing a `zero_centered` RMSNorm with a default one). Often the native XTuner ops already match HF bitwise (see §12), so no patching is needed; prefer `XTUNER_HF_IMPL` over these helpers for parity. The §8 parallel-feature tests reuse this same idiom with `cfg.dispatcher` / `cfg.ep_size` (and `FSDPConfig(ep_size=...)`) set. ### 7.4 Checkpoint paths Read from env vars; **never** hard-code. Example: `QWEN3_MOE_PATH = os.environ["QWEN3_MOE_PATH"]` at `tests/model/test_qwen3_moe.py`. Document new env vars in the PR body. ### 7.5 Drop-in training config Ship a runnable training config at `ci/config/.py` so a developer (and CI) can verify the port end-to-end with one command: load the HF checkpoint, run a few real training steps, watch the loss curve descend as expected. This is the smoke test that catches anything the unit tests miss — the full `from_hf` → `model.forward` → loss → backward → optimizer step → FSDP shard/reduce chain — and the file doubles as the example users copy when wiring the model into their own training pipeline. Mirror `ci/config/qwen3_moe_30BA3.py` (MoE) or `ci/config/qwen3_dense.py` (dense), keeping its structure: one `()` (from §3.2) fed into a `TrainerConfig` alongside `optim_cfg` / `lr_cfg` / `fsdp_cfg` / `dataset_cfg` / `dataloader_cfg` / `loss_cfg`. `load_from` and `tokenizer_path` read from an env var (§7.4 — typically the same one as the parity test). Verify by running ~50 steps and confirming the loss drops monotonically into a plausible range for that model size; record the trajectory in the PR body alongside the §6 convergence trace. --- ## 8. Training optimizations Only after the §6 baseline is committed and the §7 tests are green: verify each switch runs correctly and keeps precision within budget **against that baseline**. 1. EP and SP can be enabled normally. 2. With `ep_size > 1`, `intra_layer_micro_batch` enables correctly and stays precise. 3. `torch.compile` runs correctly. 4. fp8 can be enabled. 5. `XTUNER_ACTIVATION_OFFLOAD` can be enabled. Each switch needs a test that compares its output to the §6 baseline — e.g. the dispatcher × ep_size parity matrix (at minimum `{(None,1), ("all2all",4), ("all2all",8)}`, add `"deepep"` if supported), compiled-vs-eager, or fp8-vs-bf16 within tolerance. These items are well-suited to the multi-agent workflow in §11: the main agent analyzes and creates tasks (with test cases), sub-agents develop in parallel, and the main agent merges. Report the precision error introduced by each switch. --- ## 9. Reference cheat-sheet | Concept | Concrete location | |-------------------------|--------------------------------------------------------------------------------| | Dense base config | `xtuner/v1/model/base.py` (`TransformerConfig`) | | MoE base config | `xtuner/v1/model/moe/moe.py` (`MoEConfig`) | | Compose base config | `xtuner/v1/model/compose/base.py` (`BaseComposeConfig`) | | `from_hf` (model) | `xtuner/v1/model/base.py` (`from_hf`) | | `save_hf` (model) | `xtuner/v1/model/base.py` (`save_hf`) | | `to_hf_key_list` proto | `xtuner/v1/model/base.py` | | MoE layer split | `xtuner/v1/model/moe/moe.py` (`first_k_dense_replace`) | | RoPE auto-parse | `xtuner/v1/module/rope/rope.py` (`RopeParametersConfig.from_hf_config`)| | Dense reference | `xtuner/v1/model/dense/qwen3.py` | | MoE reference | `xtuner/v1/model/moe/qwen3.py` | | MoE w/ MLA + NoAux | `xtuner/v1/model/moe/deepseek_v3.py` | | MoE w/ layout transform | `xtuner/v1/model/moe/gpt_oss.py` | | Remote-code reference | `xtuner/v1/model/moe/qwen3.py` (`Qwen3MoEFoPEConfig`) | | Compose reference (VLM) | `xtuner/v1/model/compose/qwen3_vl/` | | Dense test reference | `tests/model/test_qwen3_dense.py` | | MoE test reference | `tests/model/test_qwen3_moe.py` | | Engine test reference | `tests/engine/test_moe_train_engine.py` | --- ## 10. Commit discipline 0. **Commit as you go — this is authorized, do not pause for a separate go-ahead.** The user invoking this skill authorizes the commits this workflow produces. Finish a logical step → lint/test → commit, then move on. Do not leave the whole port uncommitted and end with "should I commit?"; that strands the work. (Hard-to-reverse or outward actions like `git push` / opening a PR still need explicit confirmation — committing locally does not.) 1. **Follow the stacked-PR convention.** Plan upfront how many commits the model port needs — one per logical step (e.g. shared-path fixes; model class + config + registration; the `XTUNER_HF_IMPL` switch; baseline tests; then one per parallel/optimization feature). Order them by dependency (a shared-path fix the baseline relies on goes first). Land every later change as `git commit --fixup=` + `git rebase --autosquash` into the commit it belongs to. **Never grow the history with endless patch commits just to keep fixing things up.** 2. **Every commit must pass lint:** ```bash pre-commit run --files $(find xtuner/v1) ``` --- ## 11. Multi-agent workflow This applies **only to the §8 parallel/optimization features** — and only once the §6 baseline and §7 baseline tests are done and committed (§10). The baseline itself is built sequentially, not multi-agent; the §8 features are independent of each other, which is what makes them safe to develop in parallel. 1. The **main agent** splits the §8 features into tasks and, per §8, writes each feature's comparison test against the committed baseline and commits the tests. 2. **Sub-agents** branch off the baseline/test commit and develop one feature each (e.g. one switch / one tower) in parallel, making its test pass. 3. **The two sides adjust each other** — a shared-interface change is made once by the main agent (it updates the test/baseline; sub-agents rebase), and a failing test pushes the sub-agent to fix its feature, not to weaken the test. 4. The **main agent** merges, resolves conflicts, and reports the precision error each feature introduces relative to the baseline. --- ## 12. Parity debugging pitfalls Hard-won notes from aligning hybrid (GatedDeltaNet + gated-MHA) models bitwise. The meta-rule: **measure on real GPU and bisect to localize the divergence before concluding anything** — most "bugs" here are measurement artifacts, not model errors. A clean way to bisect is to feed one side's intermediate tensors into both modules and compare step by step (norm → q/k/v → qk_norm → rope → attention core → o_proj). - **Single-layer parity: the two sides' inputs must match exactly — especially the causal mask.** When you hand-build one layer's inputs, passing `attention_mask=None` to an HF attention/layer makes HF apply **no** causal mask while XTuner stays causal, producing a large *false* diff. Build a causal mask equivalent to XTuner's `seq_ctx`, or drive the layer through the model so the mask is constructed internally. - **`flash` ≠ `eager`, and HF eager upcasts softmax to fp32.** XTuner defaults to flash attention; HF reference eager does softmax in fp32. They are numerically close but not bitwise equal, so bitwise parity requires forcing eager on the XTuner side (that is exactly what `XTUNER_HF_IMPL` does). XTuner's `eager_attention` matches HF's `eager_attention_forward` bitwise. - **A bitwise forward does not imply a bitwise backward.** The eager softmax must run in **fp32** (`softmax(..., dtype=torch.float32).to(x.dtype)`, like HF), not bf16. With a bf16 softmax the forward logits can still round identically to HF (bitwise), but the *backward* differs (~1e-5 on `dL/dx`) because the gradient accumulates at bf16 precision. So always check backward bitwise separately (the `loss.backward()` test), not just forward — and align op dtypes (softmax, norms) in the backward-sensitive direction, fp32 where HF uses fp32. - **Whole-model backward is bitwise — but only with the full env stack.** Per-layer bitwise backward does not on its own imply whole-model bitwise backward; the missing pieces are `XTUNER_HF_IMPL=true` AND `XTUNER_DETERMINISTIC=true` AND the conftest must `import xtuner.v1` so the autotune-pin block in `xtuner/v1/__init__.py` runs before any test reaches fla. Without those, two distinct ULP-level sources compound through the LM chain rule and look indistinguishable from "fundamental bf16 noise": 1. **Triton autotune drift in fla.** fla's `@triton.autotune` picks tiling / num_warps / reduction-order per kernel. Across runs the choice can change, producing a ~1-ULP backward diff per linear-attention layer that the chain rule amplifies — a 32-layer LM stack turned a per-layer ~2e-3 abs diff at grad magnitude ~30 into a ~1.0 diff on the massive-activation channels of `dL/d(pixel_values)`. The pin (configs held to the first one, cache disabled) is in `xtuner/v1/__init__.py` under an `XTUNER_DETERMINISTIC` gate. It must run **before any module imports fla** — fla's `@triton.autotune` decorators evaluate at import time, so patching after the fact has no effect on kernels already constructed. That's why the patch is on the `xtuner.v1.*` import chain rather than in the conftest: `multiprocessing.spawn` children of `MultiProcessTestCase` re-import the test class top-level (which pulls `xtuner.v1.model.*`) without running conftest, so the conftest alone can't install the patch in those processes. 2. **`torch.library.custom_op` wraps for the linear-attention path** (`GatedDeltaNet`, `causal_conv1d`) produce a different backward autograd graph from the canonical fla / causal_conv1d call HF uses, even though the underlying kernel and forward output are identical. `XTUNER_HF_IMPL` routes through `get_chunk_gated_delta_rule_fn` / `get_causal_conv1d_fn` (§6.2) so the parity path matches HF byte-for-byte. Production train doesn't set either env var; the custom_op wraps + autotune still operate. These switches are **parity-test only**. If the whole-model backward is not bitwise but the per-layer backward is, suspect one of the two before suspecting depth-accumulated bf16 noise. - **`flash`(XTuner varlen) vs `flash_attention_2`(HF) is bitwise only for short, single-tile sequences.** They are the same kernel, so at a short seqlen (e.g. 32) the logits are bitwise (0.0); at longer sequences the tiling / accumulation order differs and they drift ~1e-2. So a short-seq text test can assert flash-vs-FA2 bitwise, but anything realistic-length (or a VLM whose image tokens make the sequence long) needs the **eager** path for seqlen-independent bitwise. Don't conclude "flash matches FA2" from one short input. - **Shared kernels are already bitwise; focus on what isn't.** GatedDeltaNet / linear attention go through the same `fla` / `causal_conv1d` kernels as HF, and native rms_norm / rope match HF bitwise — these come out at 0.0 with no work. Spend your effort on the path that actually differs (usually full attention). - **`output_hidden_states` off-by-one (full-forward route, < 40B only).** HF's tuple is `[emb, L0_out, …, L{n-2}_out, post_norm]`: index `i+1` is layer `i`'s pre-norm output only for `i ≤ n-2`; the **last** entry is post-norm, not the last layer's pre-norm output. Compare the last layer via `model.norm(...)` or you will see the final layer "explode". - **A big max-diff is not automatically a bug.** Late layers have *massive activations* (a few channels in the tens/hundreds); a tiny relative error there shows up as a large absolute max-diff. Look at the mean and the trend, not just the max. Conversely, feeding `randn` inputs to a single attention layer is ill-conditioned and misleads — use real (or model-produced) hidden states. - **A single-model failure is often a general-path bug.** E.g. `Dense.build_layers` not forwarding `rms_norm_type` to `DenseDecoderLayer` broke any `zero_centered` dense model, not just this one. Fix the shared invariant, not the symptom. - **Loading a standalone sub-tower uses `strict=False`.** A text tower loaded directly from a full VLM checkpoint will see vision (and deferred `mtp.*`) keys as unexpected; load with `strict=False` so they are ignored. - **Test the whole VLM, not just the text tower.** Text-tower parity leaves the vision tower and projector completely untested. A compose port needs a full-model forward-parity test with real image inputs (§7.1). The text path of the compose forward still runs a *dummy* vision forward to keep the tower in the autograd graph, so vision bugs surface even on text prompts. - **Vision attention is non-causal — give `eager_attention` a `causal=False` path.** XTuner's `eager_attention` defaulted to a causal block-diagonal mask; that is wrong for a vision tower (bidirectional within each image), which is why forcing vision to the causal eager op produced large diffs. The op now takes `causal=True|False` (False → non-causal block-diagonal mask), the vision module already calls it with `causal=False`, and the vision tower routes through `get_attn_impl_fn` so `XTUNER_HF_IMPL` selects eager. With that, the *whole* VLM is bitwise in one run, everything eager: `XTUNER_HF_IMPL` on the XTuner side (text causal-eager + vision non-causal-eager) and HF loaded with a single `attn_implementation="eager"`. Prefer eager for parity — it is seqlen-independent, unlike flash/FA2 (see the flash seqlen bullet above). - **Vision pos-embed dtype when reusing the Qwen3-VL tower.** HF's `fast_pos_embed_interpolate` returns fp32 (bilinear weights are fp32) and HF's vision forward casts it (`pos_embeds.to(hidden_states.dtype)`); XTuner's vision forward adds it without a cast. If you patch in the HF method for parity, cast the result back to the param dtype, or you get an fp32/bf16 `LayerNorm` mismatch (`expected scalar type Float but found BFloat16`). - **Hybrid-model `torch.compile`: choose `fullgraph` per layer — verify, then fall back; don't gate *whether* to compile.** When a decoder layer's forward is wrapped into a higher-order op (HOP) — e.g. by an activation-checkpoint wrapper, which `torch.compile` lowers to a HOP whose body must be functional so it can be recomputed in backward — *and* that layer is a linear-attention layer (GatedDeltaNet), the layer mutates an outer-scope variable inside the HOP (GatedDeltaNet writes back into `seq_ctx`). Under `fullgraph=True` that raises `HigherOrderOperator ... Mutating a variable not in the current scope (SideEffects)`. This is **not** a blanket "linear attention ⇒ `fullgraph=False`" rule — whether it triggers depends on whether the layer is actually HOP-wrapped and whether that linear impl really writes outer state. So for such a HOP + linear combination, **first test whether `fullgraph=True` runs**; keep it if it does (it compiles more of the layer), and only **fall back to `fullgraph=False`** if it doesn't — `fullgraph=False` survives by graph-breaking at the mutation and running that side-effecting write in eager. Make this a **per-layer** choice in `Dense.fully_shard` (decide each layer's `fullgraph`), not a decision about whether a layer is compiled at all. interface: display_name: "Add Chat Template" short_description: "Add and verify XTuner chat templates" default_prompt: "Use $add-chat-template to add or audit XTuner chat-template rendering, loss masks, and stop-token behavior for this official Hugging Face model." --- name: add-chat-template description: Add or audit an XTuner chat-template implementation and its loss mask from an official Hugging Face model repository or local model directory. Use when supporting a new model under xtuner/v1/data_proto/messages, checking an existing implementation against tokenizer or processor apply_chat_template, deciding which assistant output and stop/EOS tokens receive loss, adding regression tests, or validating that exported training behavior matches vLLM and SGLang inference. --- # Add Chat Template ## Goal Treat the official Hugging Face tokenizer, or processor for a multimodal model, as the rendering oracle. Produce the same conversation token sequence in XTuner, supervise model-generated output by default, and supervise the context-appropriate token that ends each assistant generation. Prove the behavior through XTuner's public tokenize API and, when available, the target vLLM and SGLang versions. Keep the implementation small and model-local. Do not infer a template or stop contract from another model in the family. ## 1. Establish the reference 1. Record the official HF repo or local directory, immutable revision, tokenizer class, Transformers version, and whether `trust_remote_code=True` is required. 2. Inspect the actual artifacts used by inference: `tokenizer_config.json`, `chat_template.jinja` or `chat_template`, `special_tokens_map.json`, `generation_config.json`, and `config.json`. Inspect `processor_config.json` and remote processor/template code when the official multimodal path uses `AutoProcessor.apply_chat_template`. 3. Run the bundled inventory against the same revision: ```bash python .claude/skills/add-chat-template/scripts/audit_hf_chat_template.py \ --trust-remote-code ``` Omit `--trust-remote-code` unless the official model requires it. For model-specific branches such as tools, thinking, or multimodal inputs, pass a JSON case file with `--cases `; see the script's `--help` output. The script inventories the tokenizer side; for a processor-owned multimodal template, run the corresponding `AutoProcessor` renderer separately. 4. Render every supported branch with the same `tokenizer.apply_chat_template` or `processor.apply_chat_template` entry point used by the official example. For multimodal models, do not assume tokenizer-only rendering reproduces media placeholders or processor expansion. Record all required template kwargs and their defaults, including `add_generation_prompt`, `continue_final_message`, tool definitions, thinking/reasoning flags, named templates, and multimodal options. Do not copy rendered strings from a model card when the executable tokenizer is available. 5. When the selected official template contains Jinja `{% generation %}` regions, call its `apply_chat_template` with `tokenize=True`, `return_dict=True`, and `return_assistant_tokens_mask=True`. Use its assistant mask as an additional oracle, then separately audit stop boundaries: a next-role stop token can sit outside the official generation region while still belonging to the previous assistant for SFT. 6. Pin the reference revision in tests or CI configuration. A moving HF branch is not a stable oracle. ## 2. Audit before implementing Search these integration points first: - `xtuner/v1/data_proto/templates/__init__.py` - `xtuner/v1/data_proto/messages/` - `xtuner/v1/data_proto/messages/__init__.py` - `xtuner/v1/datasets/sft_tokenize_fn/openai.py` - multimodal tokenize functions when applicable - `tests/chat_template/` and `tests/datasets/` If the model already exists, run its public tokenize path before changing it: ```python tokenize_fn = OpenaiTokenizeFunctionConfig(chat_template="").build(tokenizer) result = tokenize_fn({"messages": messages, "tools": tools}) ``` Compare the result with the official renderer and inspect existing tests. Find a concrete failing message sequence before fixing a mismatch. Do not rewrite a working implementation because its structure looks unusual. Use `HybridChatTemplate` plus `ChatMessages` only when fixed role wrappers fully express the official template. Add a dedicated `xtuner/v1/data_proto/messages/_chat.py` when rendering depends on message history, tools, reasoning, content types, or the next role. Follow `glm52_chat.py` and `qwen35_chat.py` as integration examples, not as rendering specifications for another model. ## 3. Write the stop contract first Inventory three distinct sources; do not collapse them into one `eos_token`: 1. **Configured generation stops**: `tokenizer.eos_token_id`, `config.json`, `generation_config.json`, including integer lists and stop strings. 2. **Template boundaries**: the exact token or token sequence following an assistant when the next item is user, tool/observation, system/developer, or the end of the sample. 3. **Engine stops**: the IDs and strings actually installed by the target vLLM and SGLang versions after loading the exported checkpoint. Create an explicit table before writing the mask: | Transition | Official serialized boundary | Token IDs | Engine mechanism | Loss when assistant loss is true | |---|---|---|---|---| | assistant → user | exact role/EOT sequence | exact IDs | EOS ID, stop ID/string, or parser | yes for the generated stop target | | assistant tool call → tool | exact observation boundary | exact IDs | EOS ID, stop ID/string, or parser | yes for the generated stop target | | final assistant → end | exact EOT/EOS, or documented training-only EOS | exact IDs | EOS ID or stop ID/string | yes | Verify each token by both encoding the exact boundary and converting its ID back to a token. Reject unknown-token conversions. Handle a multi-token stop as a sequence rather than assuming every stop is one special token. ### Role tokens can belong to the preceding assistant Some official templates have no dedicated per-turn EOT. The next role token is then the token the model must generate to stop the previous assistant. In that case, keep the official serialized order and assign loss only on that boundary token or sequence to the preceding assistant: ```text assistant answer user text ^ assistant loss ^ masked ``` For a GLM-like protocol this can mean supervising the user boundary after a normal answer, the observation boundary after a tool call, and the configured terminal EOS only when no following role supplies a boundary. Do not append a generic EOS after every assistant when the official inference template places a role/observation boundary there. That creates training history that vLLM and SGLang do not render. Conversely, when the official full-message renderer omits the final token that the model must generate to terminate, append exactly that one terminal token for SFT and document this single training-only suffix. Configured EOS IDs are alternatives, not a sequence to append together. Supervise the actual context-specific stop token present in each training path. ## 4. Implement rendering and loss together Prefer one cohesive renderer that returns rendered text plus a character- or token-level loss mask. Keep context-dependent boundary ownership beside the rendering branch that emits the boundary; avoid a collection of shallow helper functions. Apply these defaults: - Treat `assistant` messages with no `loss` field as `loss=True`. - Mask system/developer/user/tool inputs and role scaffolding unless the stop contract proves that a boundary is generated by the preceding assistant. - Supervise all model-generated assistant components: visible content, reasoning that remains in the rendered sample, reasoning closing syntax, tool calls, and the appropriate stop target. - Treat `loss=False` as masking the entire assistant output and its stop target. - Mask `add_generation_prompt` completely; it is prompt scaffolding, not a demonstrated model output. - Preserve explicit model behavior such as clearing historical reasoning only when the official template does so for the chosen kwargs. If character offsets are used, verify special-token and Unicode boundaries with the real tokenizer. A token overlapping a loss span needs a deliberate, tested rule. Keep a slow, independent token-level oracle for regression tests; do not call the production renderer from both sides of the same parity test. Maintain exact rendering parity before any documented training suffix: ```mermaid flowchart LR A["Official HF messages"] --> B["tokenizer.apply_chat_template"] A --> C["XTuner public tokenize API"] B --> D["Exact text and token IDs"] C --> D D --> E["Optional one terminal SFT stop"] E --> F["Loss-mask contract"] ``` Do not use decoded-text equality as the sole check: decoding can hide a token-ID mismatch. ## 5. Register the smallest integration For a dedicated message implementation, update only the required seams: 1. Export the message class from `xtuner/v1/data_proto/messages/__init__.py`. 2. Add the public name and serving stop metadata to `CHAT_TEMPLATE_MAP`. 3. Dispatch that name in `OpenaiTokenizeFunction`; update multimodal dispatch only if the model uses it. 4. Update typed CLI/config literals only where they restrict the new name. Keep `CHAT_TEMPLATE_MAP.stop_words` consistent with the discovered contract, but do not treat it as proof of engine behavior. vLLM and SGLang load their runtime stops from the exported HF artifacts and request/engine configuration. ## 6. Add good regression tests Use the real official tokenizer/processor and XTuner's public tokenize function. Mock only unavailable external services. Cover every supported template branch with the smallest useful matrix: - user → assistant, with `loss` omitted; - assistant → user in a multi-turn conversation; - assistant tool call → tool/observation → assistant; - final assistant; - `loss=False` followed by another role and at sample end; - system/developer input; - `add_generation_prompt=True`; - reasoning enabled/disabled and historical reasoning behavior; - tools and multimodal content when supported; - repeated assistant text, Unicode, and adjacent special tokens to expose faulty substring or offset alignment. Assert public behavior: 1. Official tokenizer/processor rendered text equals XTuner rendered text, except for the one explicitly documented final SFT suffix if required. 2. Official token IDs equal XTuner `input_ids` before that suffix. 3. `len(input_ids) == len(labels)` and every label is either the matching input ID or `IGNORE_INDEX`. 4. Assistant output receives loss by default; `loss=False` removes all of it. 5. Each transition's exact stop token IDs receive or do not receive loss as the stop table specifies. 6. User/tool content and generation prompts remain masked. 7. The fast implementation matches an independent slow oracle. Test token positions and IDs directly. A substring assertion can supplement but must not replace the loss assertion. Avoid snapshots that pass after both the renderer and expected string are changed to the same incorrect value. ## 7. Validate vLLM and SGLang Resolve exact target versions. Prefer the user's production versions; otherwise inspect the current installed versions and latest stable official source. For each engine: 1. Load the same exported tokenizer, `config.json`, and `generation_config.json`. 2. Compare the engine's chat-render/tokenize output with the official tokenizer/processor renderer for every stop-table transition. 3. Inspect or log the final stop token IDs/strings installed in the generation request. Distinguish tokenizer EOS, model generation-config EOS, request stop strings, and model-specific parsers. 4. Run generation smoke tests for a normal final answer and a tool call. Use a sufficiently high output limit and assert the engine reports a stop rather than a length limit, with no role leakage. Obtain the repository GPU lock before a local GPU server or generation test. If the exact runtime cannot be executed, audit the exact tagged engine source and report the result as a static source audit, not runtime compatibility. Switching from Chat Completions to Responses does not remove this requirement: for ordinary HF models both interfaces ultimately render model messages through the chat template. ## 8. Completion checklist Finish only when the report includes: - official HF repo, revision, and Transformers version; - implementation path and whether it is generic or dedicated; - rendering-parity matrix; - stop-contract table with token strings, IDs, and loss ownership; - tests proving default assistant loss and `loss=False`; - vLLM/SGLang versions and runtime or exact-tag source validation; - every intentional difference from the official renderer, normally at most a final training stop suffix. If auditing an existing implementation, state which requirements already had tests, which were missing, and the concrete mismatches found. Do not claim the implementation is correct merely because a test file exists. interface: display_name: "Check HF Config Save" short_description: "Implement and validate exported HF configs" default_prompt: "Use $check-hf-config-save to implement any missing official Hugging Face config export and validate it across Transformers and inference engines." --- name: check-hf-config-save description: Implement a missing XTuner Hugging Face config export from the official HF model when possible, then validate it against the installed Transformers round-trip and versioned inference-engine field contracts. Use when adding a model, when hf_config is missing or returns None, when changing from_hf/hf_config/save_hf, when reviewing config.json differences, or when debugging vLLM/SGLang checkpoint-load or inference failures. Also use for requests named check_hf_config_save. --- # Check HF Config Save ## Overview First ensure that a model with an official built-in Transformers config has an inverse `from_hf <-> hf_config` mapping. If that export is missing, implement it from the official HF model before using `xtuner._testing.check_hf_config_save` to test the real public `from_hf -> save_hf` path. Separate changes forced by the installed Transformers version from fields dropped by XTuner, then protect fields that exact inference engine versions use for model construction or weight loading. ## Workflow ### 1. Fix the scope and version matrix 1. Identify the source HF model directory, XTuner config class, export path, and any engine/version named by the user or failure log. 2. Record the executable environment versions for Python, Transformers, and any installed engines. Use the requested environment; otherwise follow the repo's default environment instructions. 3. A user-specified or production-log version takes precedence. If a component version is unspecified, look up the latest stable release from its official PyPI project/API or official release page and inspect the matching source tag. By default audit both vLLM and SGLang. 4. List every checked version in the result. Distinguish an installed runtime test from a static audit of an exact source tag. Use a table with at least these columns: | Component | Version | How selected | Validation | |---|---:|---|---| | Transformers | exact version | active environment | executable round-trip | | vLLM | exact version | user/log/latest official | runtime or exact-tag source audit | | SGLang | exact version | user/log/latest official | runtime or exact-tag source audit | Do not write `latest` without resolving it to an exact version and source. ### 2. Implement a missing HF config export Exercise the supported model's public conversion path first: ```python config = get_model_config_from_hf(SOURCE_HF_DIR) exported_config = config.hf_config ``` `exported_config is None`, or the corresponding public `config.save_hf(...)` raising the base missing-`hf_config` `NotImplementedError`, means the config-only export is not implemented. Classify the official model before changing code: ```mermaid flowchart TD A["Official HF model and exact revision"] --> B{"Built-in official Transformers config?"} B -->|Yes| C{"XTuner hf_config implemented?"} C -->|No| D["Implement inverse from_hf and hf_config mapping"] C -->|Yes| E["Continue round-trip validation"] D --> E B -->|No; trust_remote_code only| F["Keep hf_config as None and validate source-config copying"] ``` For a built-in official Transformers config: 1. Read the official checkpoint's raw `config.json` and resolve its exact `model_type`, `architectures`, repository revision, and official config implementation. Load it with the selected Transformers version and record the concrete `PretrainedConfig` subclass. Do not use a third-party model implementation as the reference. 2. Implement or complete `from_hf` with that official class. Map every architecture value needed by XTuner, use `getattr` for genuinely optional versioned fields, and use `RopeParametersConfig.from_hf_config` for RoPE. 3. Implement `hf_config` by constructing the same official config class from the XTuner config's current values. It must be the inverse of `from_hf`, not a cached copy of the source object. Re-emit every field consumed by `from_hf`, plus raw compatibility fields required by vLLM or SGLang even when Transformers treats them as optional or legacy. 4. Preserve the official `model_type` and architecture name. Do not substitute a generic `PretrainedConfig`, duplicate the official config class in XTuner, or return a hand-written dictionary. 5. Verify through public behavior that `config.hf_config` has the expected official type and that `config.save_hf(...)` succeeds, then continue with the helper below. If the official repository only works through `trust_remote_code` and the selected Transformers versions have no built-in config class, do not invent an `hf_config` implementation. Keep `hf_config=None`, retain the source `_hf_path`, and test the model's public HF save path that copies the original config, tokenizer, and required remote-code files. If the entire model or dispatch is not yet supported by XTuner, follow `$add_hf_model`; this skill only fills a missing exporter for an otherwise supported model. ### 3. Establish the Transformers reference Compare three raw JSON states: ```mermaid flowchart LR A["Source config.json"] --> B["Current Transformers load + save"] A --> C["XTuner from_hf + save_hf"] B --> D["Expected serialized reference"] C --> E["XTuner export"] D --> F["Public helper comparison"] E --> F ``` - Read the source `config.json` before `AutoConfig` normalization. - Load and save it directly with the active Transformers version. This is the serialized reference and captures forced defaults or `__post_init__` changes. - Build through XTuner's public `get_model_config_from_hf`/`from_hf` API and call the public `save_hf` API. - Compare the Transformers reference with the XTuner export. Do not use direct source-versus-export equality as the primary assertion: it misclassifies Transformers normalization as an XTuner bug. Inspect the exact Transformers config implementation, including generated or modular source and `__post_init__`, for every source-to-reference difference. Typical examples are derived `head_dim`, generated `layer_types`, new defaults, and serializer metadata, but never assume these examples cover a new model. ### 4. Audit inference-engine dependencies For every changed, missing, newly defaulted, or architecture-selecting field: 1. Search the exact engine tag, not an arbitrary installed or main-branch copy. 2. Follow model registration, config access, module construction, and the weight loader. Search aliases, `getattr` defaults, direct attribute access, and tensor name/shape conditions. 3. Classify the use as one of: - module/parameter registration; - checkpoint key or tensor-shape selection; - layer topology or MoE routing; - attention/RoPE behavior; - ignored or default-compatible. 4. Encode every value-sensitive dependency as `HFConfigFieldDependency`, with exact engine version, JSON pointer, expected value, reason, and an official source permalink. A field can be HF-equivalent yet engine-critical. For example, an engine may register a checkpoint parameter only when a legacy routing field has one exact value; omitting that field then becomes a real weight-loader failure. Static source inspection proves the dependency but not end-to-end engine compatibility. Run an engine checkpoint-load smoke test when the exact runtime and required hardware are available. Otherwise report `exact-tag source audit` and do not claim runtime success. Follow the repo GPU-lock instructions before any local GPU run. ### 5. Add the model regression test Delete narrow hand-written assertions that duplicate this contract, then add a model test on the real public conversion path: ```python import transformers from xtuner._testing import HFConfigFieldDependency, check_hf_config_save from xtuner.v1.model import get_model_config_from_hf def test_save_hf_matches_transformers_and_engine_contracts(): config = get_model_config_from_hf(SOURCE_HF_DIR) assert isinstance(config.hf_config, OFFICIAL_HF_CONFIG_CLASS) report = check_hf_config_save( config, SOURCE_HF_DIR, engine_dependencies=( HFConfigFieldDependency( engine="vllm", version="", path="/", expected="", reason="", source="", ), ), ) assert report.transformers_version == transformers.__version__ assert report.checked_engine_versions == ("vllm==",) ``` The helper performs two independent checks: - XTuner export matches the active Transformers direct round-trip. - Exported values satisfy the declared inference-engine contracts, even when Transformers itself does not consume those fields. Use `allowed_export_differences={"/json/pointer": "specific reason"}` only for an intentional XTuner difference that is not already an engine dependency. Every exception needs a non-empty model-level reason. ### 6. Prove the regression and run the matrix 1. Run the helper's own behavior tests. 2. Run the generated model test in the project's pinned Transformers version. 3. Run it in each user-requested/current upgraded Transformers environment. 4. If the old broken export is available, pass its `config.json` through the same helper and show that the expected missing/extra paths fail. This is the minimal proof that the test catches the original bug. 5. Run formatting and the smallest relevant model test suite. Report: - whether `hf_config` already existed, was implemented from which official config class/revision, or correctly remained `None` for remote code; - source-to-Transformers normalization paths; - Transformers-reference-to-XTuner paths (normally none, apart from documented allowed contracts); - engine dependency, expected value, exact version, and effect; - which checks were executable and which were source-only. ## Guardrails - Compare raw JSON, not only `AutoConfig` attributes; unknown compatibility fields may disappear from a newer config class while engines still read them. - Exercise public APIs and real conversion behavior. Do not mock XTuner model internals. - Do not replace the helper with a blanket list of expected keys. - Do not silently refresh an engine version in a test. Re-audit its exact source before updating the version and permalink. - HF semantic equivalence is not evidence of vLLM/SGLang compatibility. - Do not set `hf_config=None` for a model with a built-in official config merely to bypass config reconstruction or round-trip failures. - Preserve unrelated worktree changes and keep the implementation model-agnostic. ## Completion criteria The work is complete only when the official config export path exists where it should, the public round-trip matches Transformers, every audited engine contract is preserved, and the report identifies the exact official model revision and component versions used. --- name: model_normalize description: Normalize an HF model checkpoint — quantize bf16/fp16 weights to per-block FP8, and/or repack engine-named safetensors into HF-standard ~4GB shards. --- # model_normalize Two complementary capabilities for getting a training-engine checkpoint into a format that HF / vLLM / SGLang can consume: 1. **FP8 quantization** of large linears, in two modes: - **Reference-guided** — point at an existing FP8 checkpoint of the same architecture; the tool reads its index and reproduces the exact same quantization decisions on the source. - **Heuristic** — apply a built-in rule set that quantizes only the large matrix multiplications (attention, MLP, MoE experts) and keeps norms, routers, biases, embeddings, and vision tower in their original dtype. 2. **Standard repack** — re-shard arbitrary safetensors directories into the HF-standard ``model-{i:05d}-of-{n:05d}.safetensors`` layout with ~4GB shards and a fresh ``model.safetensors.index.json``. When ``to-fp8`` detects non-standard source shard names, it automatically runs the repack step on the quantized output, so the user-visible result is always a clean HF-conformant directory. ## Layout ``` .claude/skills/model_normalize/ ├── SKILL.md # this file ├── model_normalize.py # top-level CLI (to-fp8 / repack) ├── hf_to_fp8.py # FP8 conversion library + standalone CLI ├── heuristics.py # codified default quantization rule set ├── repack_hf.py # 4GB-shard repacker library + CLI └── convert_interns2_to_fp8.sh # example invocation for InternS2 ``` ## Usage ### Mode 1 — Reference-guided FP8 conversion ```bash python model_normalize.py to-fp8 \ --source /path/to/bf16-model \ --output /path/to/fp8-model \ --reference /path/to/fp8-reference ``` Behavior: - Loads the reference's ``model.safetensors.index.json``. - For each ``_scale_inv`` entry in the reference index, marks ```` as an FP8-quantized tensor. - Applies that exact set to the source as a literal name match (no regex surprise; works even when names don't follow any pattern). ### Mode 2 — Heuristic FP8 conversion (no reference) ```bash python model_normalize.py to-fp8 \ --source /path/to/bf16-model \ --output /path/to/fp8-model ``` The rule set is defined in ``heuristics.py``. Concretely, **quantize iff** the tensor name matches one of: | Group | Pattern (anchored) | | -------------------------------- | ------------------------------------------------------------------------------------------------------------------- | | Self-attention projections | ``[…]layers..self_attn.{q,k,v,o}_proj.weight`` | | Dense MLP | ``[…]layers..mlp.{gate,up,down}_proj.weight`` | | MoE per-expert (unfused) | ``[…]layers..mlp.experts..{gate,up,down}_proj.weight`` | | MoE fused-expert (3D tensors) | ``[…]layers..mlp.experts.{gate_up_proj, down_proj}`` | | MoE shared expert | ``[…]layers..mlp.shared_expert.{gate,up,down}_proj.weight`` | | Linear-attn wide projections | ``[…]layers..linear_attn.{in_proj_qkv, in_proj_z, out_proj}.weight`` | The ``[…]`` prefix allows ``model.`` and ``model.language_model.``; the same patterns repeat under ``mtp.layers..`` for multi-token-prediction blocks. **Everything else is kept in the original dtype**, in particular: - all norms (``*norm*``, ``*_layernorm.*``) - MoE routers (``mlp.gate.weight``, ``mlp.shared_expert_gate.weight``) - embeddings (``embed_tokens``, ``lm_head``, ``pos_embed``, ``patch_embed``, ``mtp.fc``) - vision tower (``model.visual.*``) - all ``*.bias`` - ``linear_attn`` control-flow tensors (``A_log``, ``conv1d.weight``, ``dt_bias``, ``in_proj_a/b.weight``, ``norm.weight``) This rule set was validated against the InternS2 / Qwen3-MoE FP8 references: it reproduces their quantization decisions on every tensor (no false positives or false negatives). ### Mode 3 — Standalone repack (no quantization) ```bash python model_normalize.py repack \ --source /path/to/messy-shards \ --output /path/to/standard-shards \ --shard-size-gb 4 ``` Greedy bin-packs tensors into ``model-{i:05d}-of-{n:05d}.safetensors`` shards of at most ``--shard-size-gb`` GiB, writes a fresh ``model.safetensors.index.json`` (including ``metadata.total_size``), and copies all non-safetensors files (config, tokenizer, etc.) through. ## When to use which mode | Situation | Mode | | ---------------------------------------------------------------------------- | ------------------ | | You already have an FP8 reference for this architecture. | Reference-guided | | New architecture / first-time conversion; standard naming. | Heuristic | | Quantized model has engine-specific shard names; need HF-conformant output. | (auto, see below) | | You only need to standardize shard naming; weights are already FP8 / bf16. | ``repack`` | The ``to-fp8`` subcommand auto-detects non-standard source shard names (anything that isn't ``model.safetensors`` or ``model-NNNNN-of-NNNNN.safetensors``) and chains the repack step. Pass ``--no-repack`` to keep the source naming. ## Quantization details - Block size **128 × 128**, dtype ``float8_e4m3fn`` with saturating cast. - Per-block scale stored as ``_scale_inv`` (``float32``) in the same shard as the parent weight. Layout matches what vLLM / SGLang load for per-block FP8. - 3D MoE-expert tensors (shape ``(E, D0, D1)``) are quantized per-expert; the per-expert scales are stacked along a leading dim. - ``config.json`` receives a ``quantization_config`` block with ``quant_method=fp8``, ``weight_block_size=[128, 128]``, ``scale_fmt=ue8m0``, and an explicit ``modules_to_not_convert`` list derived from the actual quantization decisions. ## Post-conversion handoff (mandatory for the agent) After **any** successful run of `to-fp8` or `repack`, the agent must ask the user — via `AskUserQuestion` — where to **place** the produced `*.safetensors` and `model.safetensors.index.json` files. The conversion output is treated as a staging area; the user decides the final home. Offer at least these options: 1. **Keep in place** — leave the files at the `--output` path as-is. 2. **Copy to a new directory** — let the user specify a target; copy the shards + index there. Always **copy** (`cp`), never move — the `--output` staging directory must remain intact unless the user explicitly asks to delete it as a separate step. 3. **Copy into an existing model directory** — copy shards/index into a user-specified directory while preserving its other files (tokenizer, config, etc.). If the user picks (2) or (3), execute the copy with `cp` (confirming the destination path back to the user before any overwrite). Do not infer the destination from earlier turns — the user may have a fresh target in mind each time. Cleaning up the staging directory afterwards is a separate, user-initiated step — never delete it as part of the handoff. ## Example: InternS2 Preview ```bash python model_normalize.py to-fp8 \ --source /mnt/shared-storage-user/puyudelivery/user/yehaochen/models/InternS2PreviewCandidate5 \ --output /tmp/InternS2-FP8 \ --reference /mnt/shared-storage-user/puyudelivery/user/yehaochen/models/InternS2PreviewCandidate5-FP8 ``` The source uses engine-specific names like ``model-language-0001-fused-save_rank0.safetensors``; the output will be repacked into the standard ``model-00001-of-NNNNN.safetensors`` layout. --- name: model_patch description: Merge a base HF model dir with an extra HF model dir by index diff — take tensors only present in extra, append them to a new output dir in HF-standard shard layout. Base wins on overlap. --- # model_patch Produce ``out = base ∪ (extra \ base)`` at the **tensor level**, by diffing the two ``model.safetensors.index.json`` files. Useful when ``extra`` is a sibling checkpoint (different training run, fine-tune, or fp8 conversion) that holds a few tensors the base is missing — e.g. an updated head, a newly-added MTP block, or auxiliary projections — and you want a self-contained merged dir without re-running the whole conversion pipeline. Overlap rule: **base wins**. Any tensor present in both indexes is taken from ``base``; the corresponding entry in ``extra`` is ignored. ## Layout ``` .claude/skills/model_patch/ ├── SKILL.md # this file └── model_patch.py # CLI + library ``` ## Usage ```bash python model_patch.py \ --base /path/to/base-model \ --extra /path/to/extra-model \ --out /path/to/merged-model ``` ## What it writes - ``model-{i:05d}-of-{N:05d}.safetensors`` shards (total ``N`` = base shards + newly-packed shards holding the extra-only tensors) - a fresh ``model.safetensors.index.json`` - **every non-tensor file from ``base``** — ``config.json``, tokenizer files, modeling/configuration ``.py``, ``generation_config.json``, README, etc. — copied verbatim so ``out`` is loadable on its own. Subdirectories are copied recursively. Existing entries in ``out`` are overwritten. ``extra``'s aux files are deliberately **not** copied. The skill's overlap rule is "base wins" at the tensor level, and that extends to configs: extra often carries engine-specific edits or modality keys that don't match the tensor topology we keep. If you need fields from ``extra``'s ``config.json`` (e.g. a new modality), merge them into ``out/config.json`` as a separate, explicit step. ## How shards are produced - **Base shards**: re-emitted under the new ``-of-N`` total. By default the bytes are fully copied so ``out`` is independent of ``base``. Pass ``--hardlink`` to use ``os.link`` instead (fast, but the two trees share inodes — editing one mutates the other; deleting one keeps the other intact since hardlinks are symmetric). - **Extra-only tensors**: greedily bin-packed into new shards of at most ``--shard-size-gb`` GiB (default 4), appended after the base shards. So a base with 12 shards plus 5 GB of extra-only tensors yields: ``model-00001-of-00014.safetensors`` … ``model-00012-of-00014.safetensors`` (copies of the base shards) and ``model-00013-of-00014.safetensors`` / ``model-00014-of-00014.safetensors`` (freshly written, containing only the extra tensors). ## When to use this skill | Situation | Use this? | | -------------------------------------------------------------------------------------- | --------- | | ``extra`` contains a few tensors ``base`` is missing; everything else is identical. | Yes | | You need to merge two checkpoints whose overlapping tensors **differ** in value. | No — this skill silently keeps base; you probably want a manual decision per key. | | ``extra`` is a full quantization of ``base`` and you want the quantized weights. | No — use ``model_normalize`` with ``--reference extra`` instead. | | You want a single ``model.safetensors`` instead of standard shards. | No — repack downstream. | ## Required usage protocol (mandatory for the agent) Before invoking ``model_patch.py``: 1. Confirm the three paths (``--base``, ``--extra``, ``--out``) with the user via ``AskUserQuestion``. Never infer them from earlier conversation. 2. Show the user the **diff summary** before writing, by reading the two ``model.safetensors.index.json`` files and reporting: - number of tensors in base - number of tensors in extra - number of extra-only keys that will actually be added - number of overlapping keys that will be dropped from extra 3. Ask the user to confirm those numbers before launching the run. After a successful run, **do not** ask the user where the output should go — ``--out`` is already the final destination the user named. Just report what landed there (shard count, index path, total size, and that base's aux files were mirrored over). Unlike ``model_normalize``, ``model_patch`` has no staging-area concept: there is no temp dir, and the caller already made the placement decision when they passed ``--out``. If the user wanted any of ``extra``'s aux files (e.g. a ``config.json`` that documents a new modality), call that out — those are not copied by the skill and must be merged in by hand. --- name: sphinx-debug description: > Use when `make html` or `sphinx-build` fails with `Extension error (sphinx.ext.autosummary)`, `ImportExceptionGroup`, or other Sphinx autosummary / autodoc import failures. Provides automated pdb diagnosis for `autodoc_mock_imports` issues. --- # Sphinx Debug ## Quick Diagnostic Command Run the build with maximum verbosity and pdb on exception: ```bash make html SPHINXOPTS="-vv -T -P" ``` Flags: - `-vv`: Verbose autodoc mock logging - `-T`: Full traceback - `-P`: Drop into pdb on exception ## For agents: automate pdb via monkey-patch If you are an agent and cannot interact with an interactive pdb session, **do not redesign a subprocess/pty capture system**. Instead, monkey-patch `pdb.post_mortem` so it runs scripted commands and prints the output directly to stdout. Run this one-shot script from the docs directory (`docs/en` or `docs/zh_cn`): ```bash cd docs/en # or docs/zh_cn python -c " import pdb, sys, io, traceback # Hijack pdb.post_mortem to auto-run diagnostic commands _original_post_mortem = pdb.post_mortem def scripted_post_mortem(tb=None): out = io.StringIO() p = pdb.Pdb(stdout=out, stdin=io.StringIO('')) p.use_rawinput = False p.reset() if tb is None: tb = sys.exc_info()[2] p.setup(None, tb) old_stderr = sys.stderr sys.stderr = out # List all inner exceptions p.onecmd('for i, exc in enumerate(exceptions): print(f\"[{i}] {type(exc).__name__}: {exc}\")') # Print full traceback of the first TypeError (the real root cause) p.onecmd('type_errs = [exc for exc in exceptions if type(exc).__name__ == \"TypeError\"]') p.onecmd('if type_errs: import traceback; traceback.print_exception(type(type_errs[0]), type_errs[0], type_errs[0].__traceback__)') p.do_quit('') sys.stderr = old_stderr print('=== PDB AUTOMATED OUTPUT ===') print(out.getvalue()) print('=== END OUTPUT ===') pdb.post_mortem = scripted_post_mortem from sphinx.cmd.build import build_main build_main(['-b', 'html', '-vv', '-T', '-P', '.', '_build/html']) " ``` This runs the real Sphinx build with `-P`, but when pdb drops in it automatically executes your commands and dumps the results to stdout. You read the output exactly as if you had typed the pdb commands yourself. ## What to look for in the automated output Sphinx autosummary raises `ImportExceptionGroup` when it cannot resolve a module. The **top-level message is often misleading** (e.g. `no module named xtuner.v1.ray.dataflow`). The real root cause is hidden inside the grouped exceptions. ### Common root-cause patterns 1. **TypeError about `__version__`** ``` TypeError: expected string or bytes-like object, got '__version__' ``` This happens when a mocked module (e.g. `torch`) is imported for real by an intermediate library (e.g. `fla`) that calls `packaging.version.parse()` on the mocked `__version__` object. **Fix**: Add the intermediate library to `autodoc_mock_imports` in `conf.py`. To identify the exact library, read the traceback: the frame just above `packaging/version.py` is the culprit (e.g. `fla/utils.py`). ```python autodoc_mock_imports = [ ... "fla", # or whatever library triggers the real import ] ``` 2. **Missing-comma string-concatenation bug** If `autodoc_mock_imports` contains an entry like `"scipytorchvision"`, a comma is missing between two string literals in `conf.py`: ```python # WRONG autodoc_mock_imports = [ ... "scipy" "torchvision", ... ] # CORRECT autodoc_mock_imports = [ ... "scipy", "torchvision", ... ] ``` 3. **AttributeError after the TypeErrors** ``` AttributeError: module 'xtuner.v1.ray' has no attribute 'dataflow' ``` This is usually a *secondary* failure caused by the earlier `TypeError`s. Fix the `TypeError` first, then rebuild. ## Step-by-step workflow 1. Run the automated pdb script (see "For agents" section above). 2. Read the `=== PDB AUTOMATED OUTPUT ===` block. 3. If you see `TypeError: expected string or bytes-like object, got '__version__'`, read the traceback to find the intermediate library name. 4. Add the identified library to `autodoc_mock_imports` in **both** `docs/en/conf.py` and `docs/zh_cn/conf.py`. 5. Also eyeball `autodoc_mock_imports` for missing commas while you are there. 6. Re-run `make html` (without `-P`) to verify the fix. --- name: xtuner-custom-judger description: Use when creating, modifying, or reviewing XTuner RL judgers under xtuner/v1/rl/judger or configs that build custom judgers. Enforces the preferred Judger payload contract and the BaseJudger/ComposedJudger constraints. --- # XTuner Custom Judger Use this skill whenever you implement or review an XTuner RL judger. ## Preferred Pattern Prefer subclassing `Judger`, not `BaseJudger`. For custom scoring logic, override these methods as needed: ```python class MyJudger(Judger): def preprocess(self, rollout_state: RolloutState) -> JudgerPayload: ... async def judge_payload(self, payload: JudgerPayloadBatch) -> JudgerOutputBatch: ... def postprocess(self, rollout_state: RolloutState, output: JudgerOutput) -> RolloutState: ... ``` Do not override `Judger.judge()` just to read extra fields from `RolloutState`. Put the field extraction in `preprocess()` instead. ## Why This Contract Exists `preprocess -> judge_payload -> postprocess` keeps the scoring payload small and composable: - `preprocess()` extracts only the fields required for scoring. - `judge_payload()` can run locally, in a Ray actor, or behind a judger pool. - `postprocess()` writes the output back to the original `RolloutState`. This avoids `deepcopy(RolloutState)` in composed and remote judging. Deep-copying rollout states can add serialization overhead and can be risky for large fields, Ray object refs, tensors, numpy arrays, or externally owned objects. ## BaseJudger Rule Subclass `BaseJudger` only when the implementation must own the full `judge()` / `batch_judge()` flow and cannot reasonably be expressed as a payload contract. BaseJudger-only judgers cannot be used as `ComposedJudger` branches. `ComposedJudger` intentionally calls branch `preprocess()` and `judge_payload()` rather than arbitrary `BaseJudger.judge()` implementations. Because it does not deep-copy `RolloutState`, supporting arbitrary BaseJudger-only branches would allow multiple branches to mutate the same rollout state concurrently and overwrite `reward` or other fields. ## ComposedJudger Checklist When adding a branch intended for `ComposedJudger`: 1. Ensure the class inherits `Judger`. 2. Ensure branch-specific fields are copied into the payload in `preprocess()`. 3. Ensure `judge_payload()` returns a `JudgerOutput` for one payload, or a `list[JudgerOutput]` for a batch payload. 4. Ensure `postprocess()` writes `rollout_state.reward`. 5. For multi-branch composed judging, ensure `merge_fn` creates the final reward shape, including `reward["score"]` when training expects it. ## Tests For custom judger changes, add focused tests for: - Single-sample `judge()`. - `batch_judge()` if batch support is implemented. - `ComposedJudger` routing if the judger is intended as a composed branch. - Rejection or explicit non-support if the implementation directly subclasses `BaseJudger`. --- name: xtuner-sync-supported-models description: Synchronize xtuner's supported model documentation (docs/en/pretrain_sft/advanced_tutorial/model.md and docs/zh_cn/pretrain_sft/advanced_tutorial/model.md) with the actual Config classes defined under xtuner/v1/model/. Use when (1) new TransformerConfig, MoEConfig, or BaseComposeConfig subclasses are added, removed, or renamed in xtuner/v1/model/, (2) existing model configs change their inheritance hierarchy, scale, or HuggingFace counterpart, or (3) a code review or user request points out that model.md is out of sync with the codebase. --- # Update XTuner Supported Model Docs Keep the English and Chinese `model.md` files synchronized with the actual Config classes in `xtuner/v1/model/`. ## Scan the Codebase Run the bundled scan script from the xtuner project root to discover all Config classes and their inheritance: ```bash python3 .agents/skills/xtuner-sync-supported-models/scripts/scan_model_configs.py ``` The script outputs JSON with two keys: - `configs`: list of every `*Config` class under `xtuner/v1/model/` with its parent classes and file path - `children`: parent-to-children mapping for the hierarchy tree ## What to Update Compare the script output against the two files: - `docs/en/pretrain_sft/advanced_tutorial/model.md` - `docs/zh_cn/pretrain_sft/advanced_tutorial/model.md` Both files share the same structure and must stay in sync: 1. **Base Config Classes** — configs that directly inherit from `TransformerConfig` (or `MoEConfig`) and provide a `from_hf` classmethod for loading HuggingFace weights 2. **Concrete Model Configs** — fixed-scale subclasses of the base configs above 3. **Compose Models** — multimodal configs that inherit from `BaseComposeConfig` 4. **Inheritance Hierarchy** — a text tree showing the full `XTunerBaseModelConfig` hierarchy ### Rules for the Base Config table Include these direct descendants of `TransformerConfig`/`MoEConfig`: - `Qwen2DenseConfig` - `Qwen3DenseConfig` - `DeepSeekV3Config` - `GptOssConfig` - `Qwen3MoEConfig` Exclude from the base table: - `MoEConfig` — it is an intermediate base class, not a usable model family - `Qwen3_5_VLTextMoEConfig` — it is an intermediate base with only one concrete child; its child `Qwen3_5_VLTextMoE35BA3BConfig` belongs under the MoE concrete table ### Rules for the Concrete Model table Include every concrete subclass that has fixed parameter defaults. For each row note: - `Config Class` - `Base Class / Family` - `Architecture Type`: `Dense`, `MoE`, `Dense (VL backbone)`, `MoE (VL backbone)` - `Scale / Notes`: parameter count or total/activated size; for VL backbones note "for multimodal" `DeepSeekV3Config` appears here even though it has no separate base entry (it is both base and concrete). ### Rules for the Compose Models section Include three sub-tables: 1. **Compose Base Config Classes** — `Qwen3VLBaseConfig`, `InternVLBaseConfig`, `InternS1BaseConfig` - `Qwen3VLBaseConfig`: VL model based on Qwen3 text backbone - `InternVLBaseConfig`: VL model based on InternViT + Qwen3 - `InternS1BaseConfig`: Science multimodal model based on InternViT + Qwen3 2. **Concrete Compose Model Configs** — every subclass of the above bases; for each row note the wrapped `Text Config` and scale ### Rules for the Inheritance Hierarchy tree Rebuild the tree from `XTunerBaseModelConfig` with two top-level branches: ```text XTunerBaseModelConfig ├── TransformerConfig │ ├── Dense Models │ │ ├── Qwen2DenseConfig │ │ │ └── Qwen2Dense7BConfig │ │ └── Qwen3DenseConfig │ │ ├── Qwen3Dense8BConfig │ │ │ └── Qwen3VLTextDense8BConfig │ │ ├── Qwen3Dense4BConfig │ │ │ └── Qwen3VLTextDense4BConfig │ │ └── Qwen3Dense0P6BConfig │ └── MoE Models (via MoEConfig) │ ├── DeepSeekV3Config │ ├── GptOssConfig │ │ ├── GptOss21BA3P6Config │ │ └── GptOss117BA5P8Config │ ├── Qwen3MoEConfig │ │ ├── Qwen3MoE30BA3Config │ │ │ └── Qwen3VLTextMoE30BA3Config │ │ ├── Qwen3MoE235BA22Config │ │ │ └── Qwen3VLTextMoE235BA22Config │ │ └── Qwen3MoEFoPEConfig │ └── Qwen3_5_VLTextMoEConfig │ └── Qwen3_5_VLTextMoE35BA3BConfig └── BaseComposeConfig ├── Qwen3VLBaseConfig │ ├── Qwen3VLMoE30BA3Config │ ├── Qwen3VLMoE235BA22Config │ ├── Qwen3VLDense4BConfig │ ├── Qwen3VLDense8BConfig │ └── Qwen3_5_BaseConfig │ └── Qwen3_5_VLMoE35BA3Config ├── InternVLBaseConfig │ ├── InternVL3P5Dense8BConfig │ ├── InternVL3P5MoE30BA3Config │ └── InternVL3P5Dense1BConfig └── InternS1BaseConfig ├── InternS1Config └── InternS1MiniConfig ``` When new configs are added, insert them into the appropriate branch following the same indentation style. ## Translation Notes Keep the Chinese `model.md` (`docs/zh_cn/...`) structurally identical to the English one. Translate: - Section headings - Table header cells - Description cells (e.g., "Image / Video + Text" → "图像/视频 + 文本") - Scale descriptions (e.g., "~7B parameters" → "约 7B 参数", "FoPE variant" → "FoPE 变体") Do **not** translate Config class names, file paths, or code identifiers. # Blackbox Agent Trace Reference This reference captures the XTuner black-box agent trace reconstruction workflow. It is intended for implementation and review work, not as user-facing docs. Read this reference when `xtuner-trace/SKILL.md` routes an agent trace request here. The design background lives in `docs/design/blackbox_agent_trace_reconstruction.md`. ## Contents - [Decision Flow](#decision-flow) - [Path A: Trajectory-Backed Reconstruction](#path-a-trajectory-backed-reconstruction) - [Path B: SessionServer-Only Fallback](#path-b-sessionserver-only-fallback) - [`record_synthetic_span(...)` Usage](#record_synthetic_span-usage) - [Reconstructed Span Shape](#reconstructed-span-shape) - [Viewer And Jaeger Requirements](#viewer-and-jaeger-requirements) - [Example: mini-SWE-agent Trajectory Conversion](#example-mini-swe-agent-trajectory-conversion) - [Verification](#verification) - [Common Failure Modes](#common-failure-modes) ## Decision Flow When adding trace support for a black-box agent, classify the available evidence before deciding where to instrument. 1. Confirm this is a black-box agent path: XTuner exposes an OpenAI-compatible inference endpoint to the agent through `SessionServer`; the agent runs its native loop and calls that endpoint for LLM generation. Tool execution, environment interaction, and observation handling happen inside the agent runtime and are not directly instrumented by XTuner. 2. Confirm whether the agent can save a trajectory for each run. 3. If a trajectory exists, check whether it contains tool-call metadata such as tool name, action/command, status, return code, error information, observation, or timestamp. 4. Use Path A when the trajectory contains tool-call metadata. The trajectory supplies tool semantics, and `session_server.request` spans supply the LLM request timeline. 5. Use Path B only when no useful trajectory exists. This can recover request gaps, but cannot recover reliable tool names, return codes, errors, or observations. Do not require the trajectory to contain complete LLM generate timing before using Path A. `SessionServer` remains the stable source for LLM request start/end times. If the trajectory also records LLM calls, use them as an optional consistency check. ## Path A: Trajectory-Backed Reconstruction Use this path when the agent can save a trajectory with tool-call metadata. ### Data Contract The `session_server.request` spans must provide: - `xtuner.session_id` or `xtuner.rollout_id` - span start time - span end time The trajectory should provide best-effort tool metadata: - tool event order, id, or enough context to infer order - tool name, action, or command - status, return code, or exception information - observation or output summary - tool start/end timestamps, or a single tool-related timestamp with documented semantics, if the agent records one Do not define a universal trajectory schema in trace core. Agent trajectory formats are agent-specific; keep parsing inside that agent's recipe or trajectory-to-trace adapter until multiple agents prove a shared model is worth extracting. ### Where To Add Trace Add trace at the three black-box observation boundaries. These instructions must be enough to re-create the instrumentation even when no example implementation exists in the checkout. #### 1. SessionServer Request Trace First check whether the `SessionServer` HTTP request handler already emits `session_server.request` spans. If not, add this instrumentation to the handler that receives the agent's OpenAI-compatible request. Implementation requirements: - Locate the outer request handler, such as an aiohttp/FastAPI handler named like `_handle_request(request)`. - Read enough request headers/body to determine: - `session_id` - optional rollout id - whether the request is streaming - optional trace carrier from standard headers or a request-body carrier field - Wrap the full inference request lifecycle in: ```python with trace_span( "session_server.request", attributes={ "xtuner.stage": "llm_generate", "xtuner.session_id": session_id, "xtuner.rollout_id": rollout_id or session_id, "session.stream": is_streaming, }, parent_carrier=parent_carrier, ): ... ``` - Keep the outer `session_server.request` span around prepare, backend request, response read, response post-processing, and return-to-agent. - Use `set_trace_attributes(...)` before exiting the span to record request-level status or error metadata when available. - If adding child spans, use these names: - `session_server.prepare_request` - `session_server.send_request` - `session_server.read_response` - `session_server.record_response` - Put `xtuner.session_id` on every child span too. This makes later trajectory-to-trace adapters able to filter a single agent run reliably. Do not parse agent trajectory or infer tool calls inside `SessionServer`. Its black-box responsibility is to provide the LLM request timeline. #### 2. Agent Run Boundary Trace In the black-box recipe or launcher, find the call site that invokes the agent's native run loop, such as `agent.run(...)`. Wrap only the outer run boundary: ```python with trace_span( ".run", attributes={ "xtuner.stage": "agent.run", "xtuner.session_id": session_id, "xtuner.rollout_id": rollout_id or session_id, "agent.name": agent_name, }, ): agent.run(...) ``` This span is a lifecycle boundary. Do not wrap the agent's internal model, environment, or tool methods as the default black-box trace path. LLM/tool detail should come from `SessionServer` spans plus trajectory reconstruction. #### 3. Trajectory-To-Trace Reconstruction After the native agent run finishes, call the trajectory-to-trace adapter before `close_trace()`. Prefer a `finally` block so partial trajectories can still be converted after failures: ```python return_code = 1 try: with trace_span(".run", attributes=agent_run_attributes): return_code = agent.run(...) finally: emit_reconstructed_trace( trace_jsonl=trace_jsonl_path, trajectory=trajectory_path, session_id=session_id, ) close_trace() return return_code ``` The adapter must read the trajectory, read same-session `session_server.request` spans, derive tool timing from trajectory timestamps when possible, use request gaps only for correlation or inferred fallback, and emit reconstructed synthetic spans with `record_synthetic_span(...)`. ### Adapter Responsibilities Write a trajectory-to-trace adapter for the specific agent. The adapter should: 1. Load the agent trajectory. 2. Extract ordered tool events and their metadata. 3. Load `session_server.request` spans from the same `traces.jsonl`. 4. Filter request spans by `session_id`. 5. Sort request spans by start time and assign local indexes. 6. Extract tool timestamps from the trajectory and interpret their semantics: - If the trajectory provides reliable tool start/end times, use them as the `reconstructed_tool_call.N` span boundaries and set `xtuner.inferred=false`. - If the trajectory provides only one timestamp, record what it means if known, such as tool start, tool end, observation write, or trajectory write. Use it to correlate the tool event with the surrounding `session_server.request` window. - If the trajectory timestamp semantics are unknown or incomplete, use request-gap boundaries only as an inferred fallback and set `xtuner.inferred=true`. 7. Use `session_server.request` spans as the LLM timeline and correlation windows, not as the default source of trajectory-backed tool timing. 8. When using request-gap fallback for a bounded tool event, use: ```text tool_i.start = request_i.end tool_i.end = request_{i+1}.start ``` 9. Emit `reconstructed_agent_run`, `reconstructed_llm_call.N`, and bounded `reconstructed_tool_call.N` spans with `record_synthetic_span(...)`. 10. Treat final tool events without a next LLM request as unbounded unless the trajectory has trustworthy start/end times. When request-gap fallback is used, the interval is not pure tool runtime. It can include LLM-response parsing, user confirmation time, tool execution, observation formatting, trajectory write time, and loop overhead before the next LLM request. ## Path B: SessionServer-Only Fallback Use this path only when the agent does not save a useful trajectory. ### Where To Add Trace Add trace to the `SessionServer` request path so every LLM request produces a reliable timeline. At minimum, `session_server.request` spans must include: - `xtuner.session_id` - `xtuner.rollout_id`, falling back to the session id when there is no rollout id - request start time - request end time Supporting spans such as `session_server.prepare_request`, `session_server.send_request`, `session_server.read_response`, and `session_server.record_response` may remain useful for debugging latency, but the fallback reconstruction only requires the outer `session_server.request` start/end times. Do not put tool inference inside `SessionServer`. `SessionServer` should not parse trajectories, infer tool names, track return codes, or maintain complex per-session reconstruction state. ### Fallback Reconstruction After the agent run, a fallback adapter may infer intervals from adjacent request spans: ```text inferred_tool_call_i.start = request_i.end inferred_tool_call_i.end = request_{i+1}.start ``` The fallback can emit synthetic spans such as `inferred_tool_call.N` with: ```text xtuner.stage=tool_call xtuner.synthetic=true xtuner.inferred=true xtuner.session_id= reconstruction.source=session_server_only reconstruction.method=session_server_gap ``` Do not claim unavailable metadata in this mode. SessionServer-only fallback cannot reliably recover tool name, command, return code, stderr/stdout, agent exception, skipped action, or observation content. Do not infer the final tool event without a next LLM request. The interval has no reliable end time. ## `record_synthetic_span(...)` Usage Use `record_synthetic_span(...)` to write reconstructed historical intervals back into the active XTuner trace runtime. ```python record_synthetic_span( name: str, *, start_time_unix_ns: int, end_time_unix_ns: int, attributes: Mapping[str, Any] | None = None, parent_carrier: Mapping[str, str] | None = None, status: str = "completed", error_message: str | None = None, ) -> dict[str, str] | None ``` Rules: - Call it before `close_trace()`. - Use absolute Unix nanoseconds, not monotonic timestamps. - Ensure `end_time_unix_ns > start_time_unix_ns`. - Let the API default `xtuner.synthetic=true`. - Set `xtuner.inferred=true` only when the time boundary is inferred, such as from `session_server.request` gaps. - Use `status="error"` and `error_message=...` for real errors. - Use `parent_carrier` to attach reconstructed children under one synthetic root trace. - Treat `None` return as trace-disabled no-op. Typical root and child pattern: ```python root_ids = record_synthetic_span( "reconstructed_agent_run", start_time_unix_ns=first_request.start_time_unix_ns, end_time_unix_ns=last_request.end_time_unix_ns, attributes={ "xtuner.stage": "agent.run", "xtuner.session_id": session_id, "xtuner.inferred": False, "reconstruction.method": "trajectory_session_timeline", }, ) root_carrier = None if root_ids: root_carrier = {"traceparent": f"00-{root_ids['trace_id']}-{root_ids['span_id']}-01"} record_synthetic_span( "reconstructed_tool_call.1", start_time_unix_ns=request_1.end_time_unix_ns, end_time_unix_ns=request_2.start_time_unix_ns, attributes={ "xtuner.stage": "tool_call", "xtuner.session_id": session_id, "xtuner.inferred": True, "agent.tool_call_index": 1, "agent.tool.name": "bash", "reconstruction.source": "agent_trajectory", "reconstruction.method": "session_server_gap", }, parent_carrier=root_carrier, ) ``` Do not call OpenTelemetry SDK directly from the adapter. Use XTuner trace APIs. ## Reconstructed Span Shape Trajectory-backed reconstruction should generate one synthetic root trace: ```text reconstructed_agent_run |-- reconstructed_llm_call.1 |-- reconstructed_tool_call.1 |-- reconstructed_llm_call.2 |-- reconstructed_tool_call.2 ... ``` Use these naming rules: - `reconstructed_agent_run` for the synthetic root. - `reconstructed_llm_call.N` for LLM request intervals copied from `session_server.request`. - `reconstructed_tool_call.N` for trajectory-backed tool metadata. - `inferred_tool_call.N` only for SessionServer-only fallback without trajectory metadata. Recommended root attributes: ```text xtuner.stage=agent.run xtuner.synthetic=true xtuner.inferred=false xtuner.session_id= xtuner.rollout_id= agent.name= agent.llm_call_count= agent.tool_call_count= reconstruction.source= reconstruction.method=trajectory_session_timeline ``` Recommended LLM attributes: ```text xtuner.stage=llm_generate xtuner.synthetic=true xtuner.inferred=false agent.llm_call_index=N reconstruction.source=session_server_request reconstruction.method=session_server_request_copy reconstruction.original_trace_id= reconstruction.original_span_id= reconstruction.original_operation=session_server.request ``` Recommended tool attributes: ```text xtuner.stage=tool_call xtuner.synthetic=true xtuner.inferred=true agent.tool_call_index=N agent.tool.name= agent.tool.status= agent.tool.returncode= agent.tool.command_preview= reconstruction.source= reconstruction.method=session_server_gap reconstruction.timestamp_check= ``` Do not call trajectory-backed spans `inferred_tool_call.N`. Use `reconstructed_tool_call.N` and put uncertainty in attributes. ## Viewer And Jaeger Requirements XTuner viewer and Jaeger must consume the same reconstructed spans. Do not implement one inference path in the viewer and another inference path for Jaeger. Viewer expectations: - Show one sample per agent run/session after live and completed records are merged. - Prefer the trace containing `reconstructed_agent_run` as the primary sample trace. - Keep raw `session_server.request` trace ids available for audit. - Display the reconstructed path: ```text reconstructed_agent_run -> reconstructed_llm_call.1 -> reconstructed_tool_call.1 -> ... ``` Jaeger expectations: - The viewer's Jaeger link should open the trace containing `reconstructed_agent_run`. - Opening the first raw `session_server.request` trace is a bug for reconstructed black-box samples. - A healthy reconstructed trace contains `1 + llm_count + bounded_tool_count` reconstructed spans. # XTuner Trace Patterns These patterns describe the retained trace surface in this branch: runtime, basic API, viewer, and local setup tooling. ## Basic API Use the public facade: ```python from xtuner.v1.rl.trace import ( inject_trace_context, set_trace_attributes, trace_event, trace_function, trace_span, ) ``` ### Local Span ```python with trace_span("phase.name", attributes={"xtuner.stage": "phase"}): ... set_trace_attributes({"phase.count": 3}) ``` ### Decorated Function ```python @trace_function("phase.load") def load_item(path: str) -> object: ... ``` ### Parent Carrier ```python carrier = inject_trace_context() with trace_span("child.phase", parent_carrier=carrier): ... ``` `parent_carrier` is a basic W3C context carrier. It is not tied to `RolloutState`, Ray, aiohttp, or any XTuner business object. ## End-to-End Trace Wiring Use this pattern when adding trace support to an XTuner training run. ### Launch Script Enable trace bootstrap from the launcher, guarded by `XTUNER_TRACE_ENABLED`: ```bash SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" if [ "${XTUNER_TRACE_ENABLED:-0}" = "1" ]; then source "${REPO_ROOT}/recipe/trace/setup_trace.sh" fi ``` Keep training stdout/stderr in the main training log so runtime trace messages are visible: ```bash python xtuner/v1/train/cli/rl.py \ --config "$CONFIG_PATH" \ --num-workers "$XTUNER_RL_NUM_WORKERS" \ 2>&1 | tee -a "${WORK_DIR}/training_log_${current_time}.txt" ``` ### Training Config Add `TraceConfig` to the training config and pass it to the trainer config: ```python import os from xtuner.v1.rl.trace import TraceConfig trace_config = TraceConfig( enabled=os.environ.get("XTUNER_TRACE_ENABLED") == "1", service_name="xtuner-agent-rollout", xtuner_viewer_enabled=True, xtuner_viewer_host="0.0.0.0", xtuner_viewer_port=18080, xtuner_viewer_jaeger_query_url="http://127.0.0.1:16686", enable_rollout_trace=True, ) trainer = RLColocateTrainerConfig( ..., trace_config=trace_config, ) ``` When `output_dir` is omitted in an RL trainer config, the trainer places trace artifacts under the current experiment directory, for example `work_dirs/name/20260713071356/otel/`. ### Stage Spans Before adding spans, ask the user which stages they want to observe and which metrics each stage should expose. Put fields known when entering the stage, especially fields used for grouping or filtering, in `trace_span(..., attributes=...)`. Put results known only after or during execution, such as status, duration, token counts, errors, and reward values, in `set_trace_attributes(...)`. This lets the viewer reconstruct the call chain and identify latency hotspots or abnormal stages. Infer stage: ```python import time started_at = time.monotonic() attributes = { "xtuner.stage": "infer", "xtuner.task_name": task_name, "xtuner.sample_id": sample_id, "rollout.backend": backend, } with trace_span("agent.infer", attributes=attributes): result = await run_infer(...) set_trace_attributes( { "xtuner.status": "ok" if result.ok else "error", "prompt.tokens": result.prompt_tokens, "completion.tokens": result.completion_tokens, "stage.duration_ms": int((time.monotonic() - started_at) * 1000), } ) ``` Resource acquire stage: ```python with trace_span( "sandbox.acquire", attributes={ "xtuner.stage": "acquire", "xtuner.sample_id": sample_id, "sandbox.name": sandbox_name, }, ): client = await pool.get(sandbox_name) set_trace_attributes( { "sandbox.env_id": env_id, "sandbox.image": image, "sandbox.reused": reused, } ) ``` Validation or reward stage: ```python with trace_span( "sample.validate", attributes={ "xtuner.stage": "validate", "xtuner.sample_id": sample_id, "validator.name": validator_name, }, ): record = await validate(...) set_trace_attributes( { "xtuner.status": "ok" if record.ok else "error", "reward.score": record.score, "reward.passed": record.passed, } ) if record.error is not None: set_trace_attributes( { "error": True, "error.type": record.error.type, "error.message": record.error.message, } ) ``` Do not record prompts, responses, full configs, secrets, raw headers, stack traces, or large payloads. ### Cross-Boundary Propagation When a traced operation crosses a process, actor, or HTTP boundary, keep the propagation helper in the business or transport module that owns that boundary. Do not put Ray, HTTP, rollout, agent, or judger helpers back into `xtuner/v1/rl/trace`. #### Ray Remote Calls For Ray calls, inject the current context into a plain carrier and temporarily attach that carrier to a serializable domain object or explicit call argument. ```python from collections.abc import Mapping from contextlib import contextmanager from typing import Any from xtuner.v1.rl.trace import inject_trace_context, trace_span _TRACE_CARRIER_FIELD = "_xtuner_trace_carrier" @contextmanager def attach_trace_carrier_temporarily(target: Any, carrier: Mapping[str, str]): if not carrier: yield return extra_fields = getattr(target, "extra_fields", None) if extra_fields is None: extra_fields = {} target.extra_fields = extra_fields had_previous = _TRACE_CARRIER_FIELD in extra_fields previous = extra_fields.get(_TRACE_CARRIER_FIELD) extra_fields[_TRACE_CARRIER_FIELD] = dict(carrier) try: yield finally: if had_previous: extra_fields[_TRACE_CARRIER_FIELD] = previous else: extra_fields.pop(_TRACE_CARRIER_FIELD, None) def call_ray_remote_with_trace(remote_method, *args, trace_target: Any, **kwargs): carrier: dict[str, str] = {} inject_trace_context(carrier) with attach_trace_carrier_temporarily(trace_target, carrier): return remote_method.remote(*args, **kwargs) def pop_trace_parent_carrier(target: Any) -> dict[str, str] | None: extra_fields = getattr(target, "extra_fields", None) if not isinstance(extra_fields, dict): return None carrier = extra_fields.pop(_TRACE_CARRIER_FIELD, None) if not isinstance(carrier, Mapping): return None return {str(key): str(value) for key, value in carrier.items()} ``` On the receiver side, attach the carrier to the new span: ```python parent_carrier = pop_trace_parent_carrier(trace_target) with trace_span("worker.stage", attributes=attributes, parent_carrier=parent_carrier): ... ``` Keep the helper domain-specific. Validate that exactly one trace target is used when the transport needs one target object; reject collections if restoring the carrier to the right child span would be ambiguous. #### HTTP Calls For HTTP calls, prefer W3C headers. If a third-party client or proxy may drop custom headers, also put a copied carrier into the JSON body under an internal field and remove it before forwarding the payload downstream. ```python from collections.abc import Mapping from typing import Any from xtuner.v1.rl.trace import inject_trace_context, trace_span _TRACE_CARRIER_FIELD = "_xtuner_trace_carrier" def inject_trace_carrier_into_json_body(kwargs: dict[str, Any], carrier: Mapping[str, str]) -> None: if not carrier: return payload = kwargs.get("json") if not isinstance(payload, dict): return payload = dict(payload) payload.setdefault(_TRACE_CARRIER_FIELD, dict(carrier)) kwargs["json"] = payload def extract_trace_carrier_from_mapping(payload: Mapping[str, Any] | None) -> dict[str, str] | None: if not isinstance(payload, Mapping): return None carrier = payload.get(_TRACE_CARRIER_FIELD) if not isinstance(carrier, Mapping): return None return {str(key): str(value) for key, value in carrier.items()} def remove_trace_carrier_from_mapping(payload: dict[str, Any]) -> None: payload.pop(_TRACE_CARRIER_FIELD, None) def extract_http_parent_carrier(headers: Mapping[str, Any], payload: Mapping[str, Any] | None) -> dict[str, str] | None: header_carrier = { str(key): str(value) for key, value in headers.items() if str(key).lower() in {"traceparent", "tracestate", "baggage"} } return header_carrier or extract_trace_carrier_from_mapping(payload) ``` Caller: ```python headers = dict(headers or {}) request_kwargs = {"json": payload} carrier: dict[str, str] = {} inject_trace_context(carrier) headers.update(carrier) inject_trace_carrier_into_json_body(request_kwargs, carrier) with trace_span("http.client.request", attributes={"http.method": "POST", "http.url": url}): response = await client.post(url, headers=headers, **request_kwargs) ``` Receiver: ```python headers = dict(request.headers) payload = await request.json() parent_carrier = extract_http_parent_carrier(headers, payload) with trace_span("http.server.request", attributes=attributes, parent_carrier=parent_carrier): remove_trace_carrier_from_mapping(payload) ... ``` Do not record raw headers or full payloads as span attributes. Only record derived fields such as method, route, status code, stage, IDs, and timing. ### Viewer Output When `xtuner_viewer_enabled=True`, the trace runtime logs the trace JSONL path, the XTuner viewer URL, and the equivalent manual viewer command. The viewer defaults to port `18080`; `xtuner_viewer_port=0` is still available when a run should pick a free port automatically. The viewer process inherits the training stdout/stderr; keep those runtime log lines visible in the main training log and report them to the user after the run. Useful reference points from the full trace implementation: - `examples/v1/scripts/run_rl_run.sh` - `recipe/trace/setup_trace.sh` - `examples/v1/config/agentic_rl_qwen3p5vl_mtp_ep_code.py` ## Runtime Configure tracing explicitly: ```python from xtuner.v1.rl.trace import TraceConfig, close_trace, configure_trace runtime = configure_trace( TraceConfig( enabled=True, output_dir="work_dirs/example/otel", service_name="xtuner", xtuner_viewer_enabled=True, ) ) try: ... finally: close_trace() ``` The runtime owns OTel collector setup, trace JSONL output, and optional lightweight XTuner viewer startup. ## Viewer The viewer reads Jaeger-style traces or OTel JSONL converted into Jaeger-style payloads. It should not depend on hard-coded rollout/agent span registries. Stage display rules: 1. Use `xtuner.stage` when present. 2. Else use `stage` when present. 3. Else use `stage.name` when present. 4. Else fall back to the raw span name. Useful stable attributes: - IDs: `xtuner.rollout_id`, `xtuner.group_id`, `xtuner.session_id`, `xtuner.task_name` - Status: `xtuner.status` - Stage: `xtuner.stage`, `stage`, `stage.name` - Counts/timing: `prompt.tokens`, `completion.tokens`, `http.status_code` - Errors: `error`, `error.message`, `exception.type` Avoid recording prompts, responses, full configs, secrets, raw headers, or large payloads. ## Local Setup `recipe/trace/setup_trace.sh` and the rest of `recipe/trace/` are local helper tooling for installing OTel collector binaries, starting the local Jaeger dependency, and clearing stale `recipe.trace.viewer.server` processes on `XTUNER_TRACE_VIEWER_PORT` or the default viewer port `18080`. Keep these as setup assets; do not use them to add automatic trace behavior to training configs or launch scripts unless that integration is explicitly requested. ## Rollout Starter Preset The built-in rollout starter trace is opt-in: ```python trace_config = TraceConfig( enabled=True, xtuner_viewer_enabled=True, enable_rollout_trace=True, ) ``` Its rollout-specific helpers live in `xtuner/v1/rl/trace/rollout_api.py`: - `trace_rollout_endpoint(...)` - `trace_rollout_remote(...)` Use this only for the starter rollout chain. For custom instrumentation, ask which stages and metrics the user wants before adding spans. ## Removed Legacy Semantics Do not use or recreate these removed surfaces in this branch: - `trace_remote` - `traced_rollout_endpoint` - `traced_agent_item_endpoint` - `traced_judger_endpoint` - `xtuner/v1/rl/trace/trace_utils.py` - `xtuner/v1/rl/trace/context_propagation.py` - fixed `TRACE_SPAN_*` registries - rollout/agent/judger attribute builders --- name: xtuner-trace description: Use when the user wants to instrument, inspect, or debug XTuner traces to reconstruct a sample/request execution path, follow cross-process call chains, or identify latency hotspots, bottleneck stages, and abnormal paths from span data. --- # XTuner Trace Use this skill for the current trace layer in this repository: OpenTelemetry runtime setup, the basic public API, local setup scripts, and the trace viewer. ## Current Boundary Keep this package infrastructure-only: - Runtime/configuration: `xtuner/v1/rl/trace/runtime.py` - OTel SDK adapter: `xtuner/v1/rl/trace/otel_utils.py` - Public facade: `xtuner/v1/rl/trace/__init__.py` - Basic API: `xtuner/v1/rl/trace/api.py` - Rollout starter preset: `xtuner/v1/rl/trace/rollout_api.py` - Local trace tooling and viewer: `recipe/trace/` Do not add rollout, agent, judger, Ray remote, HTTP proxy, reward, status, or session-server business semantics to `api.py`, `runtime.py`, or `otel_utils.py`. Keep rollout-specific starter behavior inside `rollout_api.py` and gated by `TraceConfig.enable_rollout_trace`. ## Viewer Choice Choose the viewer by the observation question: - Use `recipe/trace/viewer` when the user wants rollout-oriented inspection, such as comparing samples within a training step by stage duration, status, reward/filter metadata, or recorded sample call chain. - Use Jaeger when the question is not sample/step-centric, such as inspecting raw spans, service boundaries, cross-process causality, backend/client calls, collector/export behavior, or custom instrumentation that does not emit the rollout/sample attributes expected by the recipe viewer. The recipe viewer reads `traces.jsonl` and adds XTuner rollout/sample aggregation. Jaeger remains the raw trace drill-down surface for the same traced run. ## Black-Box Agent Trace Routing When the trace request involves an agent, first ask whether this is a black-box agent trace. In this mode, XTuner exposes an OpenAI-compatible inference endpoint to the agent through `SessionServer`; the agent runs its native loop and calls that endpoint for LLM generation. Tool execution, environment interaction, and observation handling happen inside the agent runtime and are not directly instrumented by XTuner. If the answer is yes, or the request mentions trajectory reconstruction, `SessionServer` request gaps, `record_synthetic_span(...)`, `reconstructed_agent_run`, `reconstructed_tool_call.N`, XTuner viewer plus Jaeger display of recovered tool calls, read [references/blackbox-agent-trace.md](references/blackbox-agent-trace.md) before planning or editing. ## Basic API These are the interfaces defined in `xtuner/v1/rl/trace/api.py` and re-exported from `xtuner.v1.rl.trace`: - `trace_span(name, attributes=None, parent_carrier=None)` - `trace_function(name=None, attributes=None)` - `trace_event(name, attributes=None)` - `set_trace_attributes(attributes)` - `inject_trace_context(carrier=None)` ## Runtime API Use these `xtuner/v1/rl/trace/runtime.py` interfaces for explicit trace runtime setup: - `TraceConfig` - `configure_trace(...)` - `close_trace()` ## Rollout Starter Preset `TraceConfig(enable_rollout_trace=True)` enables the built-in rollout starter trace. Its helpers live in `xtuner/v1/rl/trace/rollout_api.py`: - `trace_rollout_endpoint(...)` - `trace_rollout_remote(...)` Keep this layer narrow: it may depend on `RolloutState` and Ray call boundaries, but it must not leak those dependencies into the basic API. Rollout endpoint `xtuner.span_name_path` is a sample call chain, not just the current process context stack. `trace_rollout_endpoint(...)` should append the endpoint span name to an internal trace-only chain carried in `RolloutState.extra_fields`, pass that state through remote rollout boundaries, and clean the internal field at the outermost endpoint. For example, after a sample generates and then judges, `judger.run` should record a path like `single_turn_agent_loop.run -> rollout.controller.generate -> rollout.worker.generate -> judger.run` even though the local active context at judge time may only be `single_turn_agent_loop.run -> judger.run`. Keep this behavior in `rollout_api.py` and viewer tests; do not add rollout call-chain semantics to the basic trace API. ## Add Trace Workflow When adding trace instrumentation to an XTuner run: 1. Ask for or locate the launch script and training config before editing. 2. In the launch script, source `recipe/trace/setup_trace.sh` when `XTUNER_TRACE_ENABLED=1`. 3. In the training config, add `TraceConfig` and set `enabled=True`; set `xtuner_viewer_enabled=True` when the user needs interactive sample/step inspection. Set `enable_rollout_trace=True` only when the user wants the built-in rollout starter trace. 4. Before adding any `trace_span(...)` instrumentation, you must ask the user which stages they want to observe and which metrics each stage should expose. Do not infer default stages unless the user explicitly asks you to choose. 5. Add `trace_span(...)` only around the user-confirmed observed stages. Put fields known at span start in initial attributes, and update runtime or final fields with `set_trace_attributes(...)`. 6. For cross-process or request boundaries, inject a carrier with `inject_trace_context(...)` and pass it to downstream `trace_span(..., parent_carrier=carrier)`. 7. Keep transport-specific propagation at the caller boundary; do not move rollout, agent, judger, Ray, or HTTP semantics into the basic trace package. 8. Ensure the main training log includes the trace output path, XTuner viewer URL, and equivalent manual viewer command when the XTuner viewer is enabled. ## Guardrails - Do not recreate `trace_utils.py`, `context_propagation.py`, span-name registries, or business attribute builders under `xtuner/v1/rl/trace`. - Do not import `RolloutState`, agent item classes, judgers, rollout workers, Ray actors, aiohttp clients, or trainer configs from the basic trace package. - Do not call OpenTelemetry SDK directly from business code; use the basic API only when trace instrumentation is explicitly requested. - Do not record prompts, responses, full configs, secrets, raw headers, stack traces, or large payloads as attributes. - Viewer stage grouping must come from span attributes such as `xtuner.stage`, `stage`, or `stage.name`; otherwise fall back to the raw span name. For concrete patterns, read [references/trace-patterns.md](references/trace-patterns.md) before editing trace code. ## Verification Use focused checks: - `PYTHONPATH=. python -m compileall -q xtuner/v1/rl/trace recipe/trace` - `PYTHONPATH=. python -m unittest discover -s tests/rl -p 'test_trace*.py' -v` when the trace tests exist in the worktree. - `git diff --check` ### Summary This PR introduces a comprehensive GPU memory calculator for LLM training, which helps users estimate per-GPU memory consumption before actually running training jobs. This tool is essential for capacity planning and optimizing parallel configurations. ### Features - **Memory Breakdown Analysis**: Calculates detailed memory usage including: - Model parameters (embedding, attention, MLP, head) - Gradients - Optimizer states (Adam momentum & variance) - Master parameters and gradients (for mixed-precision training) - Intermediate activations - **Parallel Strategy Support**: - Data Parallelism with ZeRO Stage 1(for Megatron)/3 (FSDP for xtuenr) - Tensor Parallelism (TP) - Pipeline Parallelism (PP) and Virtual Pipeline Parallelism (VPP) for Megatron - Pipeline parallelism memory accounts for flying micro-batches during 1F1B schedule - Expert Parallelism (EP) for MoE models - **Model Architecture Support**: - Dense models (e.g., Qwen3-8B) - MoE models (e.g., Qwen3-235B-A22B, DeepSeek-V3-671B) - Multi-Head Attention (MHA) - Multi-Latent Attention (MLA) used in DeepSeek - **Training Optimizations**: - Activation checkpointing (recomputation) - Flash Attention - Chunked loss computation ### Usage ```bash python xtuner/tools/llm_calc/llm_calculator.py -c [-o ] ``` ### Example Configurations | Config File | Model | Scale | TrainFramework | |------------|-------|-------|-------| | `xtuner_qwen_dense_8b_1node.yaml` | Qwen3-8B (Dense) | 1 node | xtuner | | `xtuner_qwen_moe_30b_1node.yaml` | Qwen3-30B-A3B (MoE) | 1 node | xtuner | | `xtuner_qwen_moe_235b_64node.yaml` | Qwen3-235B-A22B (MoE) | 64 nodes | xtuner | | `xtuner_deepseekv3_671b_128node.yaml` | DeepSeek-V3-671B (MoE) | 128 nodes | xtuner | | `megatron_qwen_moe_235b_64node.yaml` | Qwen3-235B-A22B (MoE) | 64 nodes | Megatron | ### Output Example The tool outputs: 1. Total model parameters breakdown 2. Per-GPU memory consumption for different training phases (forward/backward) 3. Detailed memory breakdown by component (params, grads, optimizer states, activations) 4. Optional Excel export for further analysis ``` # python llm_calculator.py -c xtuner_qwen_dense_8b_1node.yaml ----------------------------------------------setting:---------------------------------------------- model:Qwen3-8B zero_stage:3 tp:1 pp:1 dp:8 ep:1 vpp:1 nodes:1 mbs:1 gbs:8 L:32768 recompute:True flash_attn:True capacity:1.0 LN:36 ---------------------------------------------------------------------------------------------------- Total_params_num: 7.628208160400391 B, embed_params_num: 0.57958984375 B, attn_params_num: 1.4063873291015625 B, mlp_params_num: 5.0626373291015625 B, head_params_num: 0.5795936584472656 B Max_mem for microbatch=1: 26.58784294128418 GiB. micro1_forward_last: 26.58784294128418 GiB, micro1_backward_last: 21.386390686035156 GiB Max_mem for microbatch>1: 30.401947021484375 GiB. micro1_forward_last: 26.58784294128418 GiB, micro1_backward_last: 21.386390686035156 GiB, micro2_forward_last: 30.401947021484375 GiB ---------------------------------------------------------------------memory parts--------------------------------------------------------------------- params grads opt_states master_params master_grads acts name embed 1.159180 1.159180 0.579590 0.289795 0.289795 0.000061 head 1.159187 1.159187 0.579594 0.289797 0.289797 0.625000 perlayer_attn 0.078133 0.078133 0.039066 0.019533 0.019533 1.751953 perlayer_mlp 0.281258 0.281258 0.140629 0.070314 0.070314 2.500000 all_attn 0.000000 0.000000 1.406387 0.703194 0.703194 0.000000 all_mlp 0.000000 0.000000 5.062637 2.531319 2.531319 0.000000 in_layernorm 0.000000 0.000000 0.000000 0.000000 0.000000 0.250000 ------------------------------------------------------micro1_forward_last: 26.58784294128418 GiB------------------------------------------------------ [Static OS&Param] embed, all_attn, all_mlp, head: master_params, opt_states = 11.442312240600586 GiB [Allgathered Last Layer] perlayer_attn, perlayer_mlp: params, acts = 4.6113433837890625 GiB [Allgathered Head] head: params, acts = 1.7841873168945312 GiB [Recompute Checkpoint] in_layernorm.acts * (LN - 1) = 8.75 GiB -----------------------------------------------------micro1_backward_last: 21.386390686035156 GiB----------------------------------------------------- [Static OS&Param] embed, all_attn, all_mlp, head: master_params, opt_states = 11.442312240600586 GiB [Static Grad] embed, all_attn, all_mlp, head: master_grads = 3.8141040802001953 GiB [Allgathered First Layer] perlayer_attn, perlayer_mlp: params, acts = 4.6113433837890625 GiB [Allgathered Embed] embed: params, acts = 1.15924072265625 GiB [Previous Layer Grad ReduceScatering] perlayer_attn, perlayer_mlp: grads = 0.3593902587890625 GiB [Recompute Checkpoint] None = 0 GiB -----------------------------------------------------micro2_forward_last: 30.401947021484375 GiB------------------------------------------------------ [Static OS&Param] embed, all_attn, all_mlp, head: master_params, opt_states = 11.442312240600586 GiB [Static Grad] embed, all_attn, all_mlp, head: master_grads = 3.8141040802001953 GiB [Allgathered Last Layer] perlayer_attn, perlayer_mlp: params, acts = 4.6113433837890625 GiB [Allgathered Head] head: params, acts = 1.7841873168945312 GiB [Recompute Checkpoint] in_layernorm.acts * (LN - 1) = 8.75 GiB ``` ### Technical Notes for Xtuner FSDP (with recompute) Memory Breakdown of the above example 1. At initilization, each GPU holds sharded master params and optimizer states (usually both in fp32). ``` [Static OS&Param] embed, all_attn, all_mlp, head: master_params, opt_states = 11.442312240600586 GiB ``` 2. At the end of the first micro batch forward: 2.1 The whole last decoder layer and lm head are gathered in each GPU (in bf16). ``` [Allgathered Last Layer] perlayer_attn, perlayer_mlp: params, acts = 4.6113433837890625 GiB [Allgathered Head] head: params, acts = 1.7841873168945312 GiB ``` 2.2 The checkpoint of all decoder layers also exist in each GPU ``` [Recompute Checkpoint] in_layernorm.acts * (LN - 1) = 8.75 GiB ``` 3. At the end of the first micro batch backward 3.1 Each GPU holds master params' graident now ``` [Static Grad] embed, all_attn, all_mlp, head: master_grads = 3.8141040802001953 GiB ``` 3.2 The whole first decoder layer and embedding layer are gathered in each GPU ``` [Allgathered First Layer] perlayer_attn, perlayer_mlp: params, acts = 4.6113433837890625 GiB [Allgathered Embed] embed: params, acts = 1.15924072265625 GiB ``` 3.3 FSDP does ReduceScatter of previous layer's gradients at the same time (to overlap computation and communication) ``` [Previous Layer Grad ReduceScatering] perlayer_attn, perlayer_mlp: grads = 0.3593902587890625 GiB ``` 3.4 The checkpoints for recomputing have been released, along with the forward process ``` [Recompute Checkpoint] None = 0 GiB ``` 4. At the end of the second micro batch forward, the only difference with the first micro batch is that each GPU also holds master gradients now. ``` [Static Grad] embed, all_attn, all_mlp, head: master_grads = 3.8141040802001953 GiB ``` The maximum memory usually occurs in these cases, but maybe different when under different configs such as `chunk_loss_size`.

.codex - AI Agent Protocol & Instructions

Explore Related Technologies & Topics