diff --git a/README.md b/README.md index 95659459ce8..25904be2328 100644 --- a/README.md +++ b/README.md @@ -115,6 +115,7 @@ For detailed information on backend selection, configuration, and examples, see - 🔜 **On-Policy Cross-Tokenizer Distillation** - cross-tokenizer support for the on-policy (MOPD) recipe (today MOPD is same-tokenizer; cross-tokenizer xToken is off-policy) - 🔜 **New Models** - Qwen3-Next, Minimax +- ✅ **CSD** - Cross-tokenizer distillation using bidirectional sub-token tables; see the [preparation and training guide](docs/guides/cross-tokenizer-distillation.md#csd-common-support-distillation). - ✅ **X-Token Off-Policy Distillation** - Distillation across mismatched (student, teacher) tokenizers via a precomputed projection matrix. - ✅ **Muon Optimizer** - Emerging Optimizer support for SFT/RL - ✅ **Megatron Inference** - Improved performance for Megatron Inference (avoid weight conversion). @@ -151,7 +152,7 @@ For detailed information on backend selection, configuration, and examples, see |-|-|-| |[GRPO](#grpo)|[GRPO Single Node](#grpo-single-node)|[GRPO Multi-node](#grpo-multi-node): [GRPO Qwen2.5-32B](#grpo-qwen25-32b), [GRPO Multi-Turn](#grpo-multi-turn)| |[On-policy Distillation](#on-policy-distillation)|[Distillation Single Node](#on-policy-distillation-single-node)|[Distillation Multi-node](#on-policy-distillation-multi-node)| - |[X-Token Off-Policy Distillation](#x-token-off-policy-distillation)|[X-Token Off-Policy Distillation Single Node](#x-token-off-policy-distillation-single-node)|[X-Token Off-Policy Distillation Multi-node](#x-token-off-policy-distillation-multi-node)| + |[Cross-Tokenizer Distillation](#cross-tokenizer-distillation)|[X-Token Single Node](#x-token-off-policy-distillation-single-node), [CSD Single Node](#csd-single-node)|[X-Token Multi-node](#x-token-off-policy-distillation-multi-node), [CSD Multi-node](#csd-multi-node)| |[SFT](#supervised-fine-tuning-sft)|[SFT Single Node](#sft-single-node)|[SFT Multi-node](#sft-multi-node)| |[DPO](#dpo)|[DPO Single Node](#dpo-single-node)|[DPO Multi-node](#dpo-multi-node)| |[RM](#rm)|[RM Single Node](#rm-single-node)|[RM Multi-node](#rm-multi-node)| @@ -425,12 +426,13 @@ uv run python examples/run_grpo_sliding_puzzle.py ## Distillation -NeMo RL supports two distillation recipes: +NeMo RL supports three distillation recipes: | Recipe | Multi-teacher | Asynchronous | Policy | Loss | Tokenizer | Backend | |---|---|---|---|---|---|---| | MOPD | Yes | Yes | On-policy | Top-1 sampled (RL-style) | Same | Megatron | | xToken | Yes | No (sync) | Off-policy | Full-logit (KL) | Same or different | DTensor | +| CSD | Yes | No (sync) | Off-policy | Common-support top-K (JSD) | Same or different | DTensor or Megatron | ## On-policy Distillation @@ -475,11 +477,18 @@ sbatch \ > [!NOTE] > For GB200 systems with 4 GPUs per node, use `--gres=gpu:4` instead. -## X-Token Off-Policy Distillation +## Cross-Tokenizer Distillation -We support off-policy distillation between a student and a teacher that **do not share a tokenizer** (cross-tokenizer, or "x-token", distillation) — for example, distilling a `Qwen/Qwen3-4B` teacher into a `meta-llama/Llama-3.2-1B` student. The reference recipe trains on the ungated, CC-BY-4.0 [Nemotron-Pretraining-Specialized-v1.1](https://huggingface.co/datasets/nvidia/Nemotron-Pretraining-Specialized-v1.1) corpus (`Nemotron-Pretraining-Formal-Logic` subset). +Cross-tokenizer distillation trains a student from a teacher that **does not share its tokenizer** — for example, distilling a `Qwen/Qwen3-4B` teacher into a `meta-llama/Llama-3.2-1B` student. The two models split the same text into different tokens, so their next-token distributions cover different vocabularies and cannot be compared directly. NeMo RL supports two off-policy methods that bridge this gap: -You can read about the details of the x-token distillation implementation [here](docs/guides/xtoken-off-policy-distillation.md) and [here](https://arxiv.org/abs/2605.21699), including how the (student, teacher) projection matrix is built and how the loss modes work. +1. **[X-Token](#x-token-off-policy-distillation)** projects the student's logits into the teacher's vocabulary through a precomputed (student, teacher) projection matrix. It runs on the DTensor backend. +2. **[CSD](#csd-common-support-distillation)** (Common-Support Distillation) aligns the two tokenizations into chunks and compares the models only on the token and token-chain events that both tokenizers can express, using bidirectional sub-token tables. It runs on the DTensor or Megatron backend, chosen per model. + +Both methods share the off-policy trainer, token alignment, and colocated CUDA IPC teacher transport, and both support multiple teachers, including teachers that share the student's tokenizer. Their reference recipes train on the ungated, CC-BY-4.0 [Nemotron-Pretraining-Specialized-v1.1](https://huggingface.co/datasets/nvidia/Nemotron-Pretraining-Specialized-v1.1) corpus (`Nemotron-Pretraining-Formal-Logic` subset). The [cross-tokenizer distillation guide](docs/guides/cross-tokenizer-distillation.md) covers both methods. + +### X-Token Off-Policy Distillation + +X-Token handles the vocabulary mismatch by projecting the student into the teacher's vocabulary. You can read about the details of the x-token distillation implementation [here](docs/guides/cross-tokenizer-distillation.md#x-token) and [here](https://arxiv.org/abs/2605.21699), including how the (student, teacher) projection matrix is built and how the loss modes work. Before launching a run, build the projection matrix for your (student, teacher) tokenizer pair: @@ -494,7 +503,7 @@ Before launching a run, build the projection matrix for your (student, teacher) > [!NOTE] > X-token distillation runs on the DTensor backend only, and the student and teacher must be colocated on the same node (teacher logits travel via CUDA IPC). -### X-Token Off-Policy Distillation Single Node +#### X-Token Off-Policy Distillation Single Node To run x-token distillation on a single node using `meta-llama/Llama-3.2-1B` as the student and `Qwen/Qwen3-4B` as the teacher: @@ -513,7 +522,7 @@ uv run python examples/run_xtoken_off_policy_distillation.py \ cluster.gpus_per_node=8 ``` -### X-Token Off-Policy Distillation Multi-node +#### X-Token Off-Policy Distillation Multi-node ```sh # Run from the root of NeMo RL repo @@ -535,6 +544,71 @@ sbatch \ > [!NOTE] > For GB200 systems with 4 GPUs per node, use `--gres=gpu:4` instead. +### CSD (Common-Support Distillation) + +CSD compares the student and teacher on events that both tokenizers can express — a shared token, or a shared chain of tokens — and keeps the remaining probability mass in a REST category. You can read about the details of the CSD implementation [here](docs/guides/cross-tokenizer-distillation.md#csd-common-support-distillation), including how the sub-token tables are built and how the loss treats each chunk type. + +Before launching a run, build both sub-token tables for your (student, teacher) tokenizer pair and validate them: + +```sh +uv run python -m tools.csd.build_subtoks_table \ + --source-tokenizer meta-llama/Llama-3.2-1B \ + --destination-tokenizer Qwen/Qwen3-4B \ + --output cross_tokenizer_data/csd/student_to_teacher.pt + +uv run python -m tools.csd.build_subtoks_table \ + --source-tokenizer Qwen/Qwen3-4B \ + --destination-tokenizer meta-llama/Llama-3.2-1B \ + --output cross_tokenizer_data/csd/teacher_to_student.pt + +uv run python -m tools.csd.validate_subtoks_tables \ + --student-tokenizer meta-llama/Llama-3.2-1B \ + --teacher-tokenizer Qwen/Qwen3-4B \ + --student-to-teacher cross_tokenizer_data/csd/student_to_teacher.pt \ + --teacher-to-student cross_tokenizer_data/csd/teacher_to_student.pt +``` + +> [!NOTE] +> Each CSD model runs on the DTensor or Megatron backend (the exemplar uses Megatron for both), and the student and teachers must be colocated on the same node (teacher logits travel via CUDA IPC). + +#### CSD Single Node + +To run CSD on a single node using `meta-llama/Llama-3.2-1B` as the student and `Qwen/Qwen3-4B` as the teacher: + +```sh +uv run python examples/run_csd.py +``` + +By default, this uses the configuration in `examples/configs/csd.yaml`, whose table paths point at the files built above. Tables stored elsewhere are passed per teacher (`teachers.0.aligner.student_to_teacher_table_path=...` and `teachers.0.aligner.teacher_to_student_table_path=...`). You can customize other parameters with command-line overrides. For example, to send dense teacher logits instead of the default top-128 and train for 100 steps: + +```sh +uv run python examples/run_csd.py \ + loss_fn.teacher_topk_ipc_k=0 \ + distillation.max_num_steps=100 +``` + +#### CSD Multi-node + +```sh +# Run from the root of NeMo RL repo +NUM_ACTOR_NODES=2 + +COMMAND="uv run ./examples/run_csd.py --config examples/configs/csd.yaml cluster.num_nodes=2 cluster.gpus_per_node=8 checkpointing.checkpoint_dir='results/csd_2nodes' logger.wandb_enabled=True logger.wandb.name='csd-2nodes'" \ +CONTAINER=YOUR_CONTAINER \ +MOUNTS="$PWD:$PWD" \ +sbatch \ + --nodes=${NUM_ACTOR_NODES} \ + --account=YOUR_ACCOUNT \ + --job-name=YOUR_JOBNAME \ + --partition=YOUR_PARTITION \ + --time=4:0:0 \ + --gres=gpu:8 \ + ray.sub +``` + +> [!NOTE] +> On more than one node, every teacher must use the student's data-parallel size and tensor × context × pipeline parallel group size, with each group inside one node. For GB200 systems with 4 GPUs per node, use `--gres=gpu:4` instead. + ## Supervised Fine-Tuning (SFT) We provide example SFT experiments using various datasets including [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/), OpenAI format datasets (with tool calling support), and custom JSONL datasets. For detailed documentation on supported datasets and configurations, see the [SFT documentation](docs/guides/sft.md). diff --git a/docs/design-docs/csd-lockstep-sequence-packing.md b/docs/design-docs/csd-lockstep-sequence-packing.md new file mode 100644 index 00000000000..6bd0413ed1d --- /dev/null +++ b/docs/design-docs/csd-lockstep-sequence-packing.md @@ -0,0 +1,1077 @@ +# Lockstep Sequence Packing for CSD + +Status: Design and staged qualification criteria + +Target: CSD loss and sequence-packing implementation + +Last updated: 2026-10-12 + +## Decision + +cross-tokenizer sequence packing will use one controller-owned, multi-model packing plan. +Every physical bin contains the same ordered logical samples for the student and +all teachers. Each model still tokenizes and materializes that bin independently, +so student and teacher token offsets, padding, cumulative lengths, and parallel +layouts are never shared. + +This design replaces the earlier student-only proposal. It also rejects +independently repacking each teacher: a teacher and the student may have +different token lengths, but they must not silently disagree about which logical +samples a physical bin owns. + +This is the authoritative physical-packing plan for the target branch. +Boundary-aware source segmentation is a separate preprocessing problem and is +not a substitute for the lockstep plan defined here. + +The design applies to both DTensor V2 and Megatron at the control-plane level. +The two backends require different physical packing, unpacking, and parallelism +adapters. Initial DTensor support is limited to `CP=1`; Megatron `CP>1` is an +explicit acceptance target. DTensor `CP>1` is a later milestone, but its target +layout is still boundary-aware head-tail sharding rather than a contiguous +global split of the packed token stream. + +## Why cross-tokenizer Needs a Shared Plan + +Ordinary sequence packing assumes one tokenization per sample. cross-tokenizer carries +one student tokenization and zero or more different teacher tokenizations, plus +sample-local alignment metadata. Packing independently inside each `Policy` +call is therefore insufficient: + +- a first-fit decision based on student lengths can differ from one based on + teacher lengths; +- the student and teachers can have different data-parallel layouts; +- current CUDA IPC aggregation reconstructs samples positionally; +- the generic packed-loss wrapper assumes every tensor's second dimension is + the student sequence axis; +- the CSD loss contains nonlinear scaling and optional data-parallel teacher + scoring whose call scope must not change accidentally. + +The current scalar packer in +[`BatchedDataDict.shard_by_batch_size`](../../nemo_rl/distributed/batched_data_dict.py) +is policy-local and accepts only one `input_lengths` vector. Both teacher IPC +entry points in [`Policy`](../../nemo_rl/models/policy/lm_policy.py) explicitly +reject sequence packing. This proposal adds a shared logical control plane and +keeps backend-specific tensor transformations behind `Policy`. + +## Goals + +1. Treat one original Arrow row as one logical sample. +2. Preserve a durable sample identity and a unique identity for each occurrence + in a global batch. +3. Pack the same ordered sample membership for the student and every teacher. +4. Enforce every model's real token capacity and backend padding constraints. +5. Prevent attention and next-token targets from crossing logical-sample + boundaries. +6. Return teacher outputs to the existing per-logical-sample cross-tokenizer contract. +7. Preserve CSD loss, gradient, normalization, metric, and collective semantics. +8. Support dense and sparse teacher IPC through explicit staged milestones. +9. Support Megatron context parallelism, including a required `CP=2` parity + gate, without claiming DTensor context-parallel support prematurely. + +## Non-goals + +- Character-count grouping is not part of physical sequence packing. +- This proposal does not infer sentence boundaries inside a text field. +- This proposal does not split oversized Arrow rows. Oversized rows require a + separate document-chunking policy before cross-tokenizer tokenization and alignment. +- It does not make `collator.num_packed_rows` greater than one. That collator option + is a separate data-format feature and remains `1`. +- It does not make dynamic batching compatible with sequence packing. +- It does not introduce a fused cross-tokenizer loss. The existing packed fused-loss path + supports `LOGPROB`, while cross-tokenizer requires logits and independent teacher axes. + +## Terminology + +- **Source row**: one original Arrow record. +- **Logical sample**: the unit independently tokenized, aligned, masked, and + scored. In this design it is one source row. +- **`sample_id`**: durable identity of the logical source row. +- **`batch_item_id`**: unique identity of one occurrence in one global batch. +- **Loss cohort**: the logical samples that would have formed one loss call in + the unpacked run. With the current cross-tokenizer `train_micro_batch_size=1`, every + logical sample is its own cohort. +- **Physical bin**: several logical samples concatenated for one model forward. +- **Side**: the student or one teacher. Every side has its own tokenizer, + capacity, backend, and parallelism constraints. +- **Raw cumulative lengths**: offsets over valid tokens for logical reassembly. +- **Padded cumulative lengths**: offsets over physical packed storage after + backend-required per-sample padding. + +## Dataset Contract + +### One Arrow row is one logical sample + +The target dataset contract uses the Arrow record boundary as the only known +sample/document boundary. Text inside a row is opaque; punctuation and +paragraphs do not create additional samples. + +This contract does not claim that every Arrow row is an unchunked source +document. It uses the only boundary encoded by the input schema. + +`characters_per_sample` currently accumulates complete rows until a character +threshold is reached and joins them with newlines in +[`ArrowTextDataset._pack_generator`](../../nemo_rl/data/datasets/response_datasets/arrow_text_dataset.py). +It neither slices a continuous stream into fixed character multiples nor +preserves the constituent row boundaries. Accurate row-aware sequence packing +therefore requires: + +```yaml +collator: + num_packed_rows: 1 +data: + train: + characters_per_sample: null +``` + +If a source row exceeds any model's individual context capacity, setup or +collation must fail with the offending `sample_id`. Independently truncating the +same text with different tokenizers can retain different source-character +spans, so truncation must not be used to make an overflowing sample appear +packable. + +### cross-tokenizer SFT distillation with different chat templates + +For chat or SFT data, one complete conversation row remains one logical sample. +System, user, assistant, and tool turns inside that conversation are not +physical-packing boundaries: assistant tokens must still be able to attend to +the system prompt, user prompt, and earlier turns in the same conversation. +Sample-aware cumulative lengths and position resets therefore occur between +conversations, not between messages. + +cross-tokenizer must retain the raw role/content messages until the cross-tokenizer +collator. The student and every cross-tokenizer teacher independently apply +their own chat template and tokenizer before the shared planner runs. The fully +rendered token lengths, including model-specific headers, role markers, +BOS/EOS/EOT tokens, tool scaffolding, and other template additions, are the +lengths used for capacity checks. A student-rendered string must never be reused +as a teacher's input merely because the models received the same raw messages. +Reuse is valid only when the tokenizer, chat template, and template arguments +are all identical. + +For the current cross-tokenizer path, this means using `kd_data_processor` to preserve +raw messages and `CSDCrossTokenizerCollator` with `mode="chat"`. The ordinary +single-tokenizer `sft_processor` renders and tokenizes earlier in the pipeline; +feeding its output into the cross-tokenizer chat path would either double-apply a +template or force the student's rendered scaffold onto the teachers. + +For a bin containing conversations `A` and `B`, model `m` materializes: + +```text +tokens_m: [ rendered(A, template_m) | rendered(B, template_m) ] +cu_seqlens_m: [ 0, len_m(A), len_m(A) + len_m(B) ] +position_ids_m: [ 0 .. len_m(A)-1 | 0 .. len_m(B)-1 ] +``` + +Do not conflate the packed sample base with the start of source content. For +sample `i` in one side, `cu_seqlens_m[i]` is the base of the *complete rendered +conversation*, including its BOS and template prefix. That is the attention and +position-reset boundary. Source content excluding template scaffold is not +generally one contiguous interval in a multi-turn conversation: headers and +end-of-turn markers occur between system, user, assistant, and tool content. +Represent it as per-turn semantic regions instead: + +```text +(batch_item_id, turn_index, role, region_name) + -> side-local rendered token span [start_m, end_m) + -> packed span [cu_seqlens_m[i] + start_m, + cu_seqlens_m[i] + end_m) +``` + +At minimum, retain assistant-content regions and any explicit EOT target used +by cross-tokenizer. Tool-call training additionally needs regions for the function name, +arguments, and tool result. Keep alignment metadata in the side-local coordinate +space; use the packed span only to select forward outputs, then restore the +side-local view before loss. These KD regions are separate from the student SFT +supervision mask. + +The ordered membership `[A, B]` is shared across all sides, but the cumulative +offsets generally differ because their templates and tokenizers differ. The +student `token_mask` must match the chosen unpacked SFT baseline's assistant +generation mask; `sample_mask` remains one scalar per conversation occurrence. +Teacher prompt and scaffold tokens remain part of the teacher forward context +even when the cross-tokenizer alignment and KD targets cover only semantic assistant +regions. + +The current chat aligner already finds the same assistant content inside the two +different rendered strings, constructs alignment pairs independently for each +assistant message, and translates them to positions in each side's complete +rendered conversation. Template-only role markers and prompt scaffolding remain +unaligned; when detectable on both sides, each side's possibly different EOT +token is added as an explicit paired target. Keep those pairs in +logical-sample coordinates, pack all associated fields by `batch_item_id`, and +restore the logical conversation before next-token shifting and the cross-tokenizer +loss. No packed-bin-global span or chunk-ID rebasing is required. + +The current cross-tokenizer chat collator builds the student CE mask from raw assistant +content only, while the ordinary SFT path masks the rendered assistant message +chunk. Before claiming SFT parity, make this policy explicit and compare against +the intended unpacked baseline. The recommended representation keeps two +independent products: a template-derived student SFT mask for CE, and +template-free semantic regions for cross-tokenizer KD. Model-specific role +headers need not be cross-token aligned merely because the student trains them; +semantically paired stop/EOT targets can be aligned explicitly. + +The current chat collator tokenizes with `truncation=True`. The packed path must +instead obtain exact untruncated post-template lengths for every side and fail +or invoke an explicit upstream conversation-splitting policy when one side +overflows. Tool-aware templates also require each sample's `tools` metadata to +be preserved and passed to every side's `apply_chat_template`; the current +cross-tokenizer chat collator does not yet do that. Native-thinking region alignment is +likewise a guarded follow-up rather than an implicit consequence of packing. + +Tool support is staged but must fail closed: + +1. Extend `DatumSpec` with optional `tools` and preserve a defensive copy in + `kd_data_processor`, together with the structured messages and stable sample + identity. +2. Pass the same canonical per-row tool schema to the student and every teacher + `apply_chat_template` call. Each side still uses its own configured template + and template arguments. Tool-schema/scaffold tokens count toward that side's + exact capacity and remain attention context, but are not automatically KD + targets. +3. Preserve message-level `tool_calls` and `tool` responses. Introduce semantic + target regions for assistant text, tool name, tool arguments, and EOT, then + map each region into every side's rendered token coordinates before cross-tokenizer + alignment. +4. Until step 3 is implemented and tested, reject assistant turns whose target + exists only in `tool_calls` (for example `content=None`) instead of silently + producing no CE/KD target. Supplying `tools` only as prompt context remains a + valid earlier milestone. + +### Identity + +The current collator carries `idx`, but that value is a post-processing position +and is dropped before teacher export. The packed path needs two identities: + +1. `sample_id` is created before filtering/splitting removes source provenance. + When the source has no explicit ID, derive it deterministically from the + dataset namespace and stable raw-row ordinal. Preserve an upstream UUID when + one exists. +2. `batch_item_id` is assigned centrally after global-batch validation padding + and before packing. It is unique for every occurrence, for example + `(batch_uid, canonical_slot)` encoded as an integer or fixed record. + +Validation-padding copies keep the same `sample_id`, receive distinct +`batch_item_id` values, and carry `sample_mask=0`. Routing and IPC joins use +`batch_item_id`, never `sample_id`, because legitimate repeated examples can +share a semantic identity. + +### Masks and boundaries + +For a physical bin with `N` logical samples: + +- `sample_mask` has `N` values, one per logical sample; +- the student `token_mask` remains on the student token axis; +- every cross-tokenizer teacher retains its own token mask; +- model-specific cumulative lengths provide block-diagonal attention; +- next-token shifting occurs only after restoring logical-sample boundaries. + +`sample_mask` is loss metadata. It does not isolate attention. Cumulative +lengths isolate attention, but they do not automatically mask loss numerators. +Every CE, KD, teacher-score, valid-token, valid-chunk, and metric accumulator +must explicitly respect `sample_mask`. + +## End-to-End Architecture + +```mermaid +flowchart TD + A["Global batch: logical Arrow rows"] --> B["Assign sample_id and batch_item_id"] + B --> C["Tokenize each row for student and every teacher"] + C --> D["Compute sample-local cross-tokenizer alignments"] + D --> E["Build one multi-model lockstep packing plan"] + E --> FS["Student backend materializes its token stream and cu_seqlens"] + E --> FT0["Teacher 0 materializes its token stream and cu_seqlens"] + E --> FTN["Teacher N materializes its token stream and cu_seqlens"] + FT0 --> U0["Unpack teacher outputs by batch_item_id"] + FTN --> UN["Unpack teacher outputs by batch_item_id"] + FS --> US["Restore differentiable student logical views"] + U0 --> L["cross-tokenizer logical loss adapter"] + UN --> L + US --> L + L --> O["CSD loss with original masks, denominators, and call scope"] +``` + +The physical forward is packed. The cross-tokenizer loss continues to reason about +logical samples and sample-local alignment coordinates. + +## Shared Packing Plan + +### Internal representation + +Use typed, immutable internal records rather than dictionaries with implicit +fields: + +```python +@dataclass(frozen=True) +class SidePackingPlan: + side_id: str + capacity: int + raw_lengths: tuple[int, ...] + effective_lengths: tuple[int, ...] + raw_cu_seqlens_by_bin: tuple[tuple[int, ...], ...] + padded_cu_seqlens_by_bin: tuple[tuple[int, ...], ...] + physical_tokens_by_bin: tuple[int, ...] + rank_bin_indices: tuple[tuple[int, ...], ...] + + +@dataclass(frozen=True) +class LockstepPackingPlan: + batch_uid: int + canonical_batch_item_ids: tuple[int, ...] + bins: tuple[tuple[int, ...], ...] + sides: Mapping[str, SidePackingPlan] +``` + +`bins` is authoritative and stores canonical item ordinals or IDs. Every side +must validate that it materialized exactly those IDs in exactly that order. A +side plan supplies only model-specific geometry; it cannot add, remove, or +repack logical samples. + +### Effective token cost + +For logical sample `i` and side `m`: + +```text +L[i,m] = valid tokens produced by model m's tokenizer +E[i,m] = L[i,m] after model/backend per-sequence rounding +B[m] = packed token capacity for side m +``` + +The student uses `sequence_packing.train_mb_tokens`. A teacher logits export +uses that teacher's `sequence_packing.logprob_mb_tokens`. + +A candidate bin is feasible only when every side satisfies its complete +physical-size function: + +```text +physical_size(m, candidate_items) <= B[m] +``` + +This is deliberately stronger than `sum(raw_lengths) <= capacity`. +`physical_size` includes: + +- the side's `make_sequence_length_divisible_by`; +- backend-required per-sample causal-CP head-tail padding, including Megatron's + current `2 * CP` rule and the corresponding rule chosen for a future DTensor + sample-aware CP adapter; +- Megatron TP padding when sequence parallelism is enabled; +- packed-total FP8 or HybridEP divisibility; +- fixed packed-tail sizing required by pipeline parallelism; +- any DTensor packed-total constraint required by TP. + +A same-tokenizer teacher may reuse the student's raw token IDs and lengths, but +it still receives its own `SidePackingPlan` because its backend, padding, and +capacity can differ. + +### Deterministic multi-dimensional binning + +The planner executes once on the controller after all exact token lengths are +known. A correctness-first implementation should be deterministic: + +1. Start from canonical logical-sample order within each DP shard. +2. Calculate each sample's normalized pressure + `max_m(E[i,m] / B[m])`. +3. Use a stable multi-dimensional first-fit-decreasing strategy, or an + order-preserving strategy when exact loss-call ordering requires it. +4. Place a sample only when every side's `physical_size` remains within budget. +5. Preserve one explicit ordered item list per bin. +6. Recompute and validate every side's raw and padded cumulative lengths. + +The initial implementation pins one supported algorithm. It must not reuse the +existing policy-local random or scalar packing choice under the same config +name. + +### Concrete tokenizer example + +For the first ten rows of the smoke Arrow data, the observed lengths were: + +| Row | Characters | Nano3.5 student tokens | Qwen3-30B teacher tokens | +|---:|---:|---:|---:| +| 0 | 3,226 | 585 | 575 | +| 1 | 1,919 | 388 | 382 | +| 2 | 1,655 | 323 | 322 | +| 3 | 1,265 | 375 | 431 | +| 4 | 407 | 106 | 104 | +| 5 | 2,636 | 730 | 770 | +| 6 | 5,613 | 1,226 | 1,175 | +| 7 | 1,335 | 288 | 279 | +| 8 | 2,331 | 436 | 419 | +| 9 | 1,344 | 246 | 239 | + +With a 4,096-token capacity on both sides, an order-preserving shared plan can +produce: + +```text +bin 0 IDs: [0,1,2,3,4,5,6,7] +student used: 4,021 teacher used: 4,038 + +bin 1 IDs: [8,9] +student used: 682 teacher used: 658 +``` + +The memberships are identical, but the boundaries are not: + +```text +student cu: [0,585,973,1296,1671,1777,2507,3733,4021] +teacher cu: [0,575,957,1279,1710,1814,2584,3759,4038] +``` + +In this CP=1 example with no additional backend rounding, only the bin tail is +storage padding. Megatron CP/SP and future sample-aware DTensor CP may also +pad each sample before concatenation. In every case, each row remains a +separate attention and loss sample. + +## Global Batch Size and Data Parallelism + +### GBS counts logical samples + +`train_global_batch_size=48` means 48 logical Arrow rows, not 48 physical bins. +Packing may turn those rows into 7, 8, or another number of physical forwards, +but it must not change: + +- the optimizer's global batch semantics; +- the global valid-sample count; +- the global valid-token and valid-chunk denominators; +- the number of semantic dataset examples consumed. + +For an entity running on 16 GPUs: + +```text +DP = 16 / (TP * CP * PP) +``` + +For example, `TP=4, CP=1, PP=1` gives `DP=4`, so each DP rank owns 12 of +the 48 logical samples before physical packing. `TP=4, CP=2, PP=1` gives +`DP=2`, so each DP rank owns 24 logical samples. + +### Schedule bins, not samples + +The correctness-first plan first assigns the canonical logical DP shard, then +packs within that shard. This preserves sample ownership, node-local teacher +IPC placement, and the unpacked loss schedule. Every model receives the same +logical DP assignment. + +Every rank in one entity must execute the same number of physical bins. Suppose +`DP=4` and per-rank packing initially yields `[2, 2, 2, 1]` bins: the global +count is 7. The controller deterministically splits one multi-sample bin on the +fourth shard, producing `[2, 2, 2, 2]`, or 8 global bins. Splitting is always +capacity-safe because each result is a subset of a previously feasible bin. + +If a shard has too few real logical samples to reach the required bin schedule, +the initial implementation fails setup. A later implementation may add an +explicit all-masked dummy bin, but it must never duplicate a valid loss +contribution merely to equalize forward counts. + +The split happens once in `LockstepPackingPlan`. If sample `batch_item_id=31` +is removed from a multi-sample bin to form the eighth bin, that same ID is +removed for the student and every teacher. Each side then recomputes its own +cumulative lengths. No `Policy` may independently "take a sequence out" to +repair its local schedule. + +Physical bin `k` is routed consistently to the corresponding DP rank for all +policies in the first implementation. The current multi-node CUDA IPC contract +already requires compatible, node-local student/teacher layouts, so Phase 1 +requires equal student and teacher DP. Support for differing DP degrees is a +later extension requiring ID-based redistribution and a bin count compatible +with all involved DP grids. + +## Backend Materialization + +### DTensor V2 + +The existing Automodel path flattens a microbatch to `[1, T]`, resets position +IDs, and provides packed attention metadata in +[`process_microbatch`](../../nemo_rl/models/automodel/data.py). The cross-tokenizer +adapter can reuse this representation after replacing policy-local binning with +the supplied side plan. + +Initial DTensor scope: + +- causal transformer models; +- `CP=1`; +- sequence parallelism disabled initially; +- TP=1 and TP>1 covered by parity tests; +- pipeline parallelism remains unsupported by the DTensor backend. + +DTensor currently rejects sequence packing with `CP>1` in +[`automodel/setup.py`](../../nemo_rl/models/automodel/setup.py). DTensor CP +packing is a separate later milestone requiring packing-aware CP sharding, +autograd-safe student reconstruction, raw/padded boundary plumbing, and +attention parity. It is not enabled by weakening the current guard. + +The causal-load-balancing rule itself is not Megatron-specific. DTensor V2 +enters PyTorch `context_parallel` through AutoModel's +[`create_context_parallel_ctx`](../../3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/distributed/context_parallel/utils.py). +For an ordinary causal sequence, PyTorch's default head-tail load balancer +divides the sequence into `2 * CP` chunks and assigns rank `r` chunks `r` and +`2 * CP - r - 1`. At `CP=2`: + +```text +logical order: [C0 | C1 | C2 | C3] +CP rank 0: [C0 | C3] +CP rank 1: [C1 | C2] +``` + +Pairing early, inexpensive queries with late, expensive queries balances +causal-attention work. A contiguous half split would give equal token counts +but unequal causal work. + +That existing DTensor path is not sufficient for cross-tokenizer physical packing. It +removes the explicit attention mask and treats its input as one continuous +causal sequence; it does not consume the logical-sample cumulative lengths. +Removing the setup guard for a packed stream `[A | B]` would therefore permit B +to attend to A. DTensor packed CP must instead use one of these sample-aware +implementations: + +1. wire AutoModel's Transformer Engine THD path, which consumes padded + cumulative lengths and calls `thd_get_partitioned_indices`; or +2. integrate PyTorch's `PerDocumentHeadTailLoadBalancer`/block-mask CP support. + +“Document” in that PyTorch API name means an independent causal-attention +segment. Under this design's dataset contract, that segment is exactly one +logical sample (one Arrow row). It does not imply discovering or preserving +sentence, paragraph, or nested-document boundaries inside the row. + +Either implementation must use each model's own cumulative lengths, apply the +same indices to tokens, positions, labels, masks, and student-aligned cross-tokenizer +metadata, and restore differentiable student outputs to logical-sample order. +The contiguous `capacity / CP` packed-stream split is not an accepted fallback. + +Nano3.5 requires an additional model-specific phase. The force-HF model must +actually pass packed attention boundaries to its attention kernels, and its +Mamba convolution and SSM paths must receive a sequence identifier that resets +state at every logical boundary. `cu_seqlens` alone cannot isolate recurrent +state when the model ignores it. Until output and gradient isolation tests pass, +DTensor Nano3.5 packing remains disabled. + +### Megatron + +Megatron is closer to the required physical representation. The existing +[`_pack_sequences_for_megatron`](../../nemo_rl/models/megatron/data.py) +constructs: + +- THD packed tensors; +- raw and padded cumulative lengths; +- per-sequence CP padding and CP sharding; +- `PackedSeqParams` for attention; +- `total_tokens`, which allows MCore to construct `seq_idx` for intended Mamba + state resets. + +The shared side plan must feed this existing materializer rather than duplicate +it. Megatron's effective-size calculation must include the padding rules in +`_get_pack_sequence_parameters_for_megatron`. + +The current materializer returns raw and padded cumulative lengths separately, +but its `PackedSeqParams` construction uses the padded vector for all q/kv +cumulative-length fields. cross-tokenizer postprocessing and loss preparation must not +infer raw logical lengths from that object. Phase 3 explicitly preserves both +boundary sets through teacher unpacking and differentiable student +reconstruction. + +Megatron rollout order: + +1. transformer TP=1, CP=1; +2. TP>1 and sequence parallelism; +3. CP=2 dense-IPC parity, required before declaring Megatron cross-tokenizer packing + generally available; +4. pipeline parallelism and its fixed packed-tail size; +5. larger CP values and supported FP8/HybridEP combinations; +6. Mamba/Nano only after numerical state-isolation and gradient tests. + +Megatron's existing structure indicates CP and Mamba support, but structure is +not numerical proof. Each topology is guarded until its parity test passes. + +## Teacher Forward and IPC Contract + +The least disruptive transport contract remains one teacher result per logical +sample. + +For each teacher bin: + +1. Materialize the packed teacher stream from its `SidePackingPlan`. +2. Run one packed forward. +3. Slice the output with that teacher's padded physical offsets and raw valid + lengths. +4. Store one per-logical-sample dense or sparse result in persistent IPC + storage. +5. Attach `batch_item_id` and the existing TP/CP shard metadata. +6. Aggregate TP/CP shards by `batch_item_id`. +7. Canonicalize results to `LockstepPackingPlan.canonical_batch_item_ids` and + assert that every real item occurs exactly once. + +[`aggregate_per_sample_handles`](../../nemo_rl/models/policy/utils.py) currently +assumes sorted DP rank plus local position recreates global order. The packed +path must use explicit IDs instead. + +### Dense IPC + +Both DTensor and Megatron full-logit postprocessors currently reject packing. +Their packed implementations should mirror the existing top-k unpacking paths: + +- use the teacher's own boundaries; +- preserve local vocabulary shards under TP; +- restore the teacher's raw valid length; +- invert the backend-native head-tail/THD permutation to logical order under + CP, then emit contiguous per-sample IPC slices where the existing schema + requires them; +- publish stable per-sample handles only after IPC storage is fully allocated. + +Variable bin cardinality must not trigger buffer growth after handles have been +published, because reallocation invalidates prior CUDA IPC handles. Preallocate +from the plan's maximum required storage or use a stable per-sample slab. + +### Sparse IPC + +DTensor sparse packing must unpack values, token IDs, `logZ`, forced-token +membership, and any support masks with exactly the same teacher boundaries. + +Megatron sparse IPC supports unpacked batches only. Its non-IPC packed top-k path is +a useful TP/CP reference, but cross-tokenizer additionally needs globally correct top-k +IDs, `logZ`, forced IDs, and the current sparse handle schema. Sparse Megatron +support is a separate milestone after dense parity. + +## Student Forward and cross-tokenizer Loss Adapter + +The generic +[`SequencePackingLossWrapper`](../../nemo_rl/algorithms/loss/wrapper.py) cannot +be used unchanged. It truncates every rank-two-or-higher tensor with the student +length, while cross-tokenizer carries independent student-token, teacher-token, and +alignment-pair axes. + +A cross-tokenizer adapter must: + +1. Slice differentiable student logits using student raw/padded boundaries. +2. Restore logical rows in canonical `batch_item_id` order. +3. Retrieve each teacher's result by `batch_item_id`. +4. Preserve each teacher's `T_teacher` axis and alignment's pair axis. +5. Keep alignment chunk IDs and spans in sample-local coordinates. +6. Apply next-token shifting inside each logical sample. +7. Invoke the CSD loss at its original logical loss-call scope. +8. Aggregate metrics by semantic numerator/denominator, not by summing row + means. + +Any vocabulary preselection that scans a restored `[B, T]` rectangle must mask +padding before choosing columns. Zero-filled padding reconstructed after a +packed forward is not guaranteed to reproduce the logits that an unpacked +padded forward would have produced. + +No bin-global alignment rebasing is necessary because packed coordinates are +removed before cross-tokenizer preparation. This also prevents a final token in sample A +from predicting the first token in sample B. + +### CSD loss-call scope + +Packing is loss-equivalent only if the loss-call scope is preserved: + +- CE and CSD chunk numerators remain additive when divided by the original + global valid-token and valid-chunk denominators. +- `dynamic_loss_scaling` computes a detached CE/KD ratio per loss call. +- dynamic teacher weights and `select_teacher` perform data-parallel scoring + collectives. + +The current cross-tokenizer configuration uses logical `train_micro_batch_size=1`. +Phase 1 restores one CSD invocation per original logical sample and retains the +same logical call count on every DP rank. Physical binning changes forwards, not +the semantic loss calls. + +For `train_micro_batch_size>1`, the implementation must either keep each +original loss cohort atomic and call CSD once per cohort, or refactor CSD into +additive statistics followed by one nonlinear combine at the original scope. +Until that exists, packing with dynamic loss scaling requires logical MBS=1. + +Initial support is restricted to collective-safe modes: + +- `kd_loss_mode=sum` with static weights (`sum_weights_metric=null`); +- dynamic loss scaling only with logical MBS=1. + +`averaged_logits`, `select_teacher`, and dynamic `sum_weights_metric` remain +guarded until their packed/unpacked parity tests pass. The latter two also need +an identical fixed collective schedule and loss-call cohort on every DP rank. +Full support must not rely on the fact that a particular binning happened to +produce equal row counts. + +### Global normalization + +`global_valid_seqs`, `global_valid_toks`, and per-teacher +`global_valid_chunks_by_idx` are computed from the original logical DP batch +before packed microbatch iteration. The same values are passed unchanged to +every logical loss call. They are never recomputed from packed tail padding or +the number of physical bins. + +Existing DTensor and Megatron CP replication/scaling conventions remain +unchanged. If an adapter changes which tensors are replicated, its denominator +scaling must be re-derived and covered by the corresponding CP parity test. + +Masked validation copies contribute to no numerator or denominator. The CSD +chunk path must gate `pair_valid` and all derived chunk counts with +`sample_mask`, not only gate the precomputed denominator. Metrics such as +`num_valid_samples` must use `sample_mask.sum()` rather than the rectangular +batch dimension. + +## Configuration and Setup Guards + +The final feature reuses the existing per-policy `sequence_packing` blocks. No +new `collator.num_packed_rows` interpretation is introduced. + +When cross-tokenizer packing is enabled: + +1. Student and every teacher must enable sequence packing; mixed enabled and + disabled policies are rejected. +2. Dynamic batching must be disabled for every side. +3. `characters_per_sample` must be null for the row-as-sample dataset contract. +4. `collator.num_packed_rows` must remain `1`. +5. Exact, pre-truncation token lengths must be available for every side. +6. Every single logical sample must fit every side's individual context and + packed capacity. +7. Every teacher must define the packed export budget used by the planner + (`sequence_packing.logprob_mb_tokens`). +8. The shared algorithm is selected once by cross-tokenizer; policies may not select + different algorithms independently. +9. Multi-node CUDA IPC initially requires equal student/teacher DP and the + existing node-local compatible model-parallel layout. +10. DTensor requires `CP=1`; sequence parallelism is disabled initially. +11. Megatron enables only topology combinations whose parity gates have passed. +12. `sequence_packing.fuse_loss` remains false for cross-tokenizer. +13. Unsupported nonlinear teacher modes or logical MBS values fail during + setup, before workers or clusters are created. + +The implementation must fail loudly rather than silently fall back to +independent packing or untracked positional routing. + +## Backend Support Matrix + +| Capability | DTensor V2 | Megatron | +|---|---|---| +| Shared logical plan and IDs | Required common code | Required common code | +| Transformer packed forward | Existing base path, adapt to supplied plan | Existing THD path, adapt to supplied plan | +| TP | Test TP=1 and TP>1 | Test TP=1 then TP>1 | +| Sequence parallelism | Initially off; validate later for supported transformers | Staged after TP | +| CP=1 | Initial target | Initial target | +| CP>1 | Deferred sample-aware head-tail/THD work; keep current guard | In scope; CP=2 is a required acceptance gate | +| PP | Backend remains PP=1 | Staged after dense CP parity | +| Dense teacher IPC | Packed unpack/routing required | Packed unpack/routing required | +| Sparse teacher IPC | Packed unpack/routing required | New sparse producer/transport required | +| Nano3.5 attention | Boundary plumbing and parity required | Packed attention path intended; prove parity | +| Nano3.5 Mamba | Explicit conv/SSM reset plumbing required | `seq_idx` path intended; prove numerical isolation | +| cross-tokenizer loss adapter | Required | Required | +| Fused packed loss | Unsupported | Unsupported | + +The architecture therefore holds for both backends, but it is not one physical +implementation and it is not initially CP-complete on DTensor. + +## Implementation Plan + +### Phase 0: Reference semantics and setup guards + +- Disable character grouping in a dedicated cross-tokenizer packing recipe. +- Preserve raw structured chat messages, optional row-level `tools`, and the + source identity through cross-tokenizer preprocessing. +- Define the unpacked cross-tokenizer SFT reference semantics: a student-template SFT + mask for CE and template-free, sample-local semantic regions for KD. +- Reject tool-call-only assistant targets until their semantic region mapper is + implemented; do not silently turn them into zero-target samples. +- Preserve or synthesize `sample_id` before dataset split/map operations. +- Add `batch_item_id` after validation padding. +- Add setup checks for uniform packing, oversize rows, dynamic batching, MBS, + nonlinear loss modes, backend topology, and IPC locality. +- Capture packed-off loss, gradients, metrics, logits, IDs, and denominators as + deterministic reference fixtures. + +Exit criterion: the packed path is still disabled, and every unsupported +configuration fails before cluster creation with a precise message. + +### Phase 1: Shared planner and plan-driven sharding + +- Add immutable `LockstepPackingPlan` and `SidePackingPlan` records under the + data-packing package. +- Add deterministic multi-dimensional feasibility and packing. +- Render each chat side with its own template, template arguments, tokenizer, + and row-level tools before planning; retain per-turn semantic region maps and + reject side-specific overflow without truncation. +- Preserve canonical DP ownership and equalize physical-bin counts by shared + bin splitting. +- Add `BatchedDataDict.shard_by_packing_plan` or an equivalent explicit API. +- Pass side plans into teacher IPC calls and student `Policy.train`. +- Reject any worker-side attempt to recompute bin membership. + +Exit criterion: CPU tests prove all capacities, identical bin membership/order, +the 7-to-8 split, deterministic routing, and ID coverage. + +### Phase 2: DTensor transformer, CP=1 + +- Materialize supplied student and teacher side plans using the Automodel packed + representation. +- Add dense packed teacher unpacking and ID-keyed IPC aggregation. +- Add the cross-tokenizer student loss adapter. +- Validate TP=1 and TP>1 with sequence parallelism off. +- Add sparse packed teacher IPC after dense loss/gradient parity. + +Exit criterion: transformer packed/unpacked logit, loss, gradient, metric, and +attention-isolation parity at CP=1 for dense and then sparse teacher export. + +### Phase 3: Megatron transformer, including CP=2 + +- Feed shared side plans into the existing THD materializer. +- Preserve raw and padded cumulative lengths separately through postprocessing + and loss preparation. +- Implement dense TP/CP teacher unpacking and ID-keyed IPC. +- Validate TP=1, TP>1, sequence parallelism, then CP=2. +- Add PP only after CP dense parity. + +Exit criterion: packed/unpacked loss and gradient parity at CP=1 and CP=2; +changing one logical sample cannot change another sample's logits. + +### Phase 4: Backend-specific sparse and hybrid-model support + +- Implement Megatron sparse IPC, including values, IDs, `logZ`, forced IDs, and + CP routing. +- Thread packed boundaries through DTensor Nano3.5 attention. +- Thread sequence IDs through DTensor Nano3.5 convolution and SSM paths. +- Numerically validate Megatron Mamba state resets. + +Exit criterion: dense and sparse cross-tokenizer parity plus attention and recurrent-state +isolation for every declared model/backend pair. + +### Phase 5: Extended parallelism + +- Add Megatron PP and larger CP topology gates. +- Add supported FP8/HybridEP combinations. +- Implement DTensor CP packing with per-logical-sample head-tail partitioning, + using TE THD or PyTorch sample-aware CP. Apply one permutation consistently + to tokens, positions, labels, masks, and cross-tokenizer student-axis metadata, and + implement its inverse for output reconstruction. Retain the existing guard + until output, loss, gradient, and isolation parity pass. +- Generalize ID routing to differing student/teacher DP only if transport and + collective semantics support it. + +Exit criterion: each enabled topology has a committed recipe and an end-to-end +nightly test. Untested matrix entries remain rejected. + +## Test Plan + +### Pure planner tests + +- every bin contains the same ordered IDs for all sides; +- every side's complete physical-size function stays within capacity; +- raw and padded cumulative lengths are monotonic and end at the expected size; +- one sample that exceeds any side fails with its `sample_id`; +- deterministic output across repeated runs; +- GBS 48, DP 4, seven initial bins becomes eight by a shared split; +- validation duplicates retain semantic ID, receive unique occurrence IDs, and + remain masked; +- same-tokenizer sides share raw lengths but retain independent physical plans. +- chat capacity uses each side's complete post-template length, including + special/template tokens, and rejects side-specific overflow without + truncation. +- tool schemas reach every side's chat template and change that side's exact + length without changing shared bin ownership. + +### Routing and IPC tests + +- arbitrary row and bin permutations reconstruct canonical order by + `batch_item_id`; +- TP/CP shards for one item assemble exactly once; +- missing, duplicated, stale, or wrong-bin IDs fail; +- dense packed output equals independent per-sample output; +- sparse values, IDs, `logZ`, forced support, and masks equal unpacked output; +- persistent IPC buffer handles remain valid for the complete step. + +### Attention and state isolation + +- perturbing sample A does not change sample B logits; +- packed logits equal independent forwards over every valid token; +- position IDs restart at every logical boundary; +- chat turns within one conversation retain causal visibility, while tokens in + another packed conversation remain invisible; +- Megatron attention isolation passes at CP=1 and CP=2; +- deferred DTensor CP=2 tests assert per-sample head-tail membership and inverse + permutation before comparing packed/unpacked logits and gradients; +- Mamba convolution and SSM state reset at every boundary; +- the same tests compare parameter gradients after the cross-tokenizer loss. + +### Loss and metric parity + +- asymmetric `T_student != T_teacher != P` axes survive unpacking; +- packed/unpacked CE and CSD KD loss parity; +- packed/unpacked parameter-gradient parity; +- chat/SFT parity with different student and teacher templates, the intended + unpacked student SFT mask, sample-local semantic KD regions, and multiple + assistant turns; +- tool-schema-as-context parity and a fail-closed test for tool-call-only + assistant targets until semantic tool regions are supported; +- static multi-teacher sum and same-vocab/direct-KL combinations; +- averaged-logits mode where supported; +- dynamic loss scaling parity at logical MBS=1; +- negative setup tests for unsupported MBS and nonlinear teacher modes; +- same-vocabulary top-k or vocabulary-column selection ignores restored padding; +- global valid-token/chunk denominators are invariant to packing; +- `sample_mask=0` removes every CE, KD, scoring, count, and metric contribution; +- metric numerators/counts yield the same reported values independent of bin + cardinality; +- equal collective call order across DP ranks. + +### Backend topology matrix + +- DTensor transformer: TP=1 and TP>1, CP=1; +- DTensor future CP=2: validate per-sample head-tail indices, raw/padded + boundaries, attention isolation, output reconstruction, and gradient parity + before removing the guard; +- Megatron transformer: TP=1, TP>1/SP, CP=1, CP=2, then PP; +- dense IPC before sparse IPC for every new backend topology; +- mixed backend student/teacher combinations only after each individual side + passes and node-local IPC constraints are satisfied; +- Nano/Mamba cases only after dedicated state-isolation tests. + +## End-to-End Slurm Qualification Gate + +The feature is not qualified by unit and parity tests alone. After those tests +pass, run two sequential Slurm smokes with the same production-size model pair: + +- student: `meta-llama/Llama-3.1-8B`; +- teacher: `Qwen/Qwen3-14B`; +- 2 nodes x 8 GPUs, GBS 128, logical MBS 1, maximum sequence length 4,096; +- TP=2, CP=2, PP=1 for both sides; +- sequence packing enabled for the student and teacher, with 4,096-token packed + train/export budgets unless measured memory requires a smaller explicitly + documented common budget; +- dynamic batching disabled; +- `characters_per_sample=null` and `collator.num_packed_rows=1`; +- static CSD teacher aggregation and logical MBS=1 unless a later phase has + separately qualified broader loss modes. + +Preflight the exact deterministic sample stream needed by both runs against +every side's untruncated effective length. A raw-text corpus may contain +rows that exceed a 4,096-token context. Do not hide them with tokenizer +truncation or re-enable character grouping. Use one explicit, reproducible +upstream row-splitting or all-side-fit filtering policy, record affected +`sample_id` values, and use the identical resulting stream for the five- and +100-step gates. + +The primary CP=2 qualification route is Megatron because CP=2 is a required +Megatron acceptance target in this design. Do not substitute DTensor CP=2 unless +its separate sample-aware CP milestone and parity suite have already passed. + +### Scheduler preflight + +Select a GPU allocation on the deployment's scheduler that supports the model +pair, requested topology, and run duration. Check node availability and use +`sbatch --test-only` when available to validate the resource request before +submitting the qualification run. Record the allocation and environment with +the validation results so another deployment can reproduce the setup. + +### Five-step gate and repair loop + +Run five optimizer steps first with checkpointing and external logging disabled +unless they are themselves under test. A pass requires all of the following: + +1. the resolved configuration proves packing is enabled on both sides, TP=2, + CP=2, character grouping is null, collator-level row packing is one, and + dynamic batching is disabled; +2. runtime packing telemetry proves that logical samples were actually combined + into physical bins—the enabled flag alone is insufficient; +3. the log reaches `Step 5/5` and `Max steps reached` with finite loss, CE, KD, + and gradient-related values; +4. no traceback, CUDA OOM, device assertion, actor death, fatal NCCL error, + attention-boundary violation, NaN/Inf, or IPC-handle failure occurs; +5. Slurm/Ray teardown finishes cleanly. + +Do not trust the parent Slurm state alone: some Ray launch paths can return zero +after a driver failure. Validate the driver log and semantic completion markers. +If the run fails, wait for the job and Ray workers to terminate, diagnose the +first causal failure, make the smallest correctness-preserving code/config/ +launcher fix, run relevant unit tests, and submit a new uniquely named five-step +attempt. Never edit the shared checkout while an attempt is still running and +never repeatedly resubmit an unchanged failure. + +### One-hundred-step gate + +Only after a five-step attempt satisfies every criterion, run 100 steps with the +same model pair, topology, data, loss settings, and packing settings. Change only +the step/scheduler/checkpoint/logging fields required for the longer run. Save a +mid-run checkpoint when practical and require a finalized `step_100` checkpoint. +A pass requires all 100 finite metric records, semantic max-step completion, a +valid finalized checkpoint, no cross-sample leakage/fatal runtime errors, and +resolved-config plus packing-telemetry artifacts. Apply the same diagnose, +fix, unit-test, and uniquely named rerun loop until it passes or a genuine +external infrastructure blocker is demonstrated. + +Record every job ID, partition, exact command/override set, commit and dirty-diff +identity, resolved config, driver/Slurm log paths, checkpoint path, result, and +failure diagnosis in the session experiment ledger. + +## Observability + +Add controller-side metrics that are independent of backend implementation: + +- logical samples per global batch; +- physical bins per side and per DP rank; +- raw and padded tokens per side; +- bin utilization and tail-padding waste per side; +- splits performed only to equalize the DP schedule; +- masked logical samples; +- dense/sparse IPC bytes and reconstruction fallback counts. + +Never report physical-bin count as sample count. Packing efficiency should be +reported separately for each side because identical membership can have very +different utilization under different tokenizers. + +## Risks and Mitigations + +### Shared bins can reduce one side's utilization + +The joint constraint may create more bins than independently optimal packing. +Correct sample ownership is the priority. Measure utilization per side before +considering more complex routing. + +### Long rows can be hidden by tokenizer truncation + +The current collator uses `truncation=True`. Compute exact lengths before the +packing decision and fail or invoke an explicit upstream chunking policy. Do not +silently accept independently truncated tails. + +### Nonlinear loss modes can change semantics + +Guard unsupported modes until the adapter preserves original cohorts and a +fixed collective schedule. Do not call the complete CSD loss once per physical +bin. + +### Packed teacher buffers can invalidate handles + +Allocate stable storage before publishing handles and validate buffer generation +metadata during consumption. + +### Hybrid models need more than attention boundaries + +Mamba convolution and SSM state require explicit resets. Keep Nano/Mamba packing +disabled until perturbation and gradient tests prove isolation. + +### Backend support can be overstated + +Capability is topology-specific. Setup guards and recipes are removed only for +matrix entries with end-to-end parity coverage. + +## Acceptance Criteria + +The first generally usable release is complete when: + +1. one Arrow row remains one logical sample with durable identity; +2. GBS and all global normalization semantics are unchanged by packing; +3. one authoritative plan controls student and teacher bin membership; +4. the 48-row/7-bin/DP schedule case is deterministic and shared; +5. attention, next-token labels, and hybrid-model state cannot cross samples; +6. teacher dense IPC and the cross-tokenizer student adapter pass packed/unpacked loss and + gradient parity; +7. DTensor transformer CP=1 passes TP coverage; +8. Megatron transformer passes CP=1 and CP=2 dense parity; +9. unsupported sparse, nonlinear-loss, Mamba, PP, or DTensor-CP combinations + fail during setup rather than silently falling back; +10. every enabled backend/topology has a maintained recipe and test. +11. chat/SFT uses exact per-side rendered lengths and preserves semantic content + regions independently of model-specific template scaffold; +12. row-level tools reach every side's template, while unsupported tool-call + target alignment fails explicitly rather than silently dropping loss. +13. the packed Qwen3-14B to Llama-3.1-8B TP2/CP2 five-step gate passes, followed + by a clean 100-step run with a finalized `step_100` checkpoint. + +Sparse transport, Nano/Mamba, PP, differing DP, and DTensor CP become supported +only when their own later-phase exit criteria pass. diff --git a/docs/guides/cross-tokenizer-distillation.md b/docs/guides/cross-tokenizer-distillation.md new file mode 100644 index 00000000000..9791735db75 --- /dev/null +++ b/docs/guides/cross-tokenizer-distillation.md @@ -0,0 +1,764 @@ +# Cross-Tokenizer Distillation + +NeMo RL supports off-policy distillation between a student and a teacher that **do not share a tokenizer** — for example, distilling a Qwen3-4B teacher into a Llama-3.2-1B student. The two models split the same text into different tokens, so their next-token distributions cover different vocabularies and cannot be compared directly. + +Tokenizer vocabularies overlap only partially. The table below reports the pairwise overlap (intersection divided by the smaller vocabulary, on canonical token forms) across several model tokenizers; the off-diagonal entries sit well below `1.0` — that shared-vocabulary gap is what cross-tokenizer distillation bridges. + +| Model | Mistral-NeMo-Minitron-8B | Qwen3-8B-Base | Llama-3.2-1B | gemma-3-4b-it | OLMo2-8B-SuperBPE-t160k | gpt-oss-20b | +|---|---:|---:|---:|---:|---:|---:| +| Mistral-NeMo-Minitron-8B | 1.0000 | 0.5119 | 0.5525 | 0.7430 | 0.4215 | 0.7591 | +| Qwen3-8B-Base | 0.5119 | 1.0000 | 0.8481 | 0.6320 | 0.4103 | 0.6462 | +| Llama-3.2-1B | 0.5525 | 0.8481 | 1.0000 | 0.6739 | 0.4980 | 0.7977 | +| gemma-3-4b-it | 0.7430 | 0.6320 | 0.6739 | 1.0000 | 0.4918 | 0.6545 | +| OLMo2-8B-SuperBPE-t160k | 0.4215 | 0.4103 | 0.4980 | 0.4918 | 1.0000 | 0.3657 | +| gpt-oss-20b | 0.7591 | 0.6462 | 0.7977 | 0.6545 | 0.3657 | 1.0000 | + +NeMo RL provides two methods: + +1. **[X-Token](#x-token)** routes student logits through a precomputed **projection matrix** that maps each student token to the teacher tokens it most plausibly corresponds to, projecting the student into the teacher's vocabulary so the two distributions can be compared. +2. **[CSD](#csd-common-support-distillation)** (Common-Support Distillation) aligns the student and teacher tokenizations of the same text into chunks and compares the two models only on events that both tokenizers can express — a shared token or a shared chain of tokens — using bidirectional **sub-token tables**. + +| | X-Token | CSD | +|---|---|---| +| Vocabulary artifact | Weighted projection matrix (`tools/x_token/`) | Student-to-teacher and teacher-to-student sub-token tables (`tools/csd/`) | +| Compared distribution | Student projected into the teacher vocabulary | Shared token and token-chain events, plus REST | +| Training entrypoint | `examples/run_xtoken_off_policy_distillation.py` | `examples/run_csd.py` | +| Exemplar configs | `xtoken_off_policy_distillation.yaml`, `xtoken_multiteacher_off_policy_distillation.yaml` | `csd.yaml`, `csd_multiteacher.yaml` | +| Training backends | Automodel | Automodel or Megatron, chosen per model | +| Teacher logits | Dense, over CUDA IPC | Sparse top-K (default `K=128`) or dense, over CUDA IPC | + +Both methods share the off-policy trainer, the token aligner, the cross-tokenizer collator, and the colocated CUDA IPC teacher transport. The top-level `algorithm` field (`xtoken` or `csd`) selects the method, and each entrypoint loads its own exemplar. + +## X-Token + +X-Token handles the vocabulary mismatch by routing student logits through a precomputed projection matrix. This section explains how to: + +1. Create the projection matrix from a (student, teacher) tokenizer pair. +2. Launch a distillation with the projection matrix. + +### How it works + +A full run has two phases. The three prep steps are *offline data prep* — +small CLI tools you run once per (student, teacher) pair — and the result is a +single `.pt` file. The final step is the actual distillation training loop. + +``` + ┌──────────────────────────────────────────────┐ + │ Offline projection-matrix preparation │ + │ │ + │ ┌────────────────────────────────────┐ │ + (student, teacher) │ │ 1. minimal_projection_via_ │ │ + tokenizers ────▶│ │ multitoken.py │ │ + │ │ — multi-token mappings │ │ + │ └─────────────────┬──────────────────┘ │ + │ │ │ + │ ┌─────────────────▼──────────────────┐ │ + │ │ 2. reapply_exact_map.py │ │ + │ │ — pin exact 1-to-1 matches │ │ + │ └─────────────────┬──────────────────┘ │ + │ │ │ + │ ┌─────────────────▼──────────────────┐ │ + │ │ 3. sort_and_cut_projection_matrix │ │ + │ │ .py — trim to runtime top_k │ │ + │ └─────────────────┬──────────────────┘ │ + └────────────────────│─────────────────────────┘ + │ + ▼ projection_matrix.pt + ┌────────────────────────────────────────────────────┐ + │ 4. examples/ │ + │ run_xtoken_off_policy_distillation.py │ + │ — align student & teacher tokens, then │ + │ teacher forward + student forward, │ + │ then x-token KD loss │ + └────────────────────────────────────────────────────┘ +``` + +The projection matrix is a sparse `[V_student, top_k]` tensor that the +training-time loss multiplies against the student logits to project them into +the teacher's vocab space. + +Projection-matrix weights mapping student tokens (rows) to teacher tokens (columns) + +Each row of the matrix holds the weights `W_{s,t}` that distribute a student +token `s ∈ V_S` over the teacher tokens `t ∈ V_T` it corresponds to. Tokens +shared by both vocabularies map 1-to-1 (e.g., `_the`, `_cat`, `_run`), while a +student token that the teacher splits into pieces spreads its weight across +those pieces (e.g., `201` → `2`, `0`, `1`). Rows are trimmed to the runtime +`top_k` in Step 3, so low-weight tail entries are dropped (hatched cell). + +### Quickstart — single command + +For the typical case, `tools/x_token/build_projection_matrix.sh` chains +the prep steps with auto-derived intermediate paths: + +```bash +./tools/x_token/build_projection_matrix.sh \ + --student-model meta-llama/Llama-3.2-1B \ + --teacher-model Qwen/Qwen3-4B \ + --runtime-top-k 4 +``` + +The wrapper writes the final matrix to +`cross_tokenizer_data/projection_matrix___top.pt` +(override with `--final-output`). Further tweaks to step 1 defaults can be configured using `--no-{scale-trick,reverse-pass,special-token-mapping}`. Run `./tools/x_token/build_projection_matrix.sh +--help` for the full list of options. + +The per-step recipes below are for advanced customization (non-default +weight thresholds, hand-picked intermediate filenames, etc.). + +### Backend and scope + +- **Automodel only.** Set `policy.automodel_cfg.enabled=true`. The Megatron policy + worker is not wired for cross-tokenizer distillation. +- **Teacher logits travel via CUDA IPC**, so student and teacher policies must + be colocated on the same node. No remote-Ray transport for x-token logits. +- **One colocated GPU pool, serial execution.** All teachers and the student + share a single `RayVirtualCluster` + (`max_colocated_worker_groups = len(teachers) + 1`). There is no separate + generation worker (off-policy distillation trains on a fixed dataset) and no + colocation toggle. Each step time-slices the shared GPUs: one teacher at a + time is onloaded, runs its forward pass, ships its logits to the student over + CUDA IPC, then offloads; the student trains last, reading and releasing each + teacher's IPC buffer. GPU memory therefore holds exactly one teacher's params + *or* the student's at any instant, plus the teachers' resident IPC logit + buffers. + +``` + RayVirtualCluster — one GPU pool, shared by every worker group + max_colocated_worker_groups = len(teachers) + 1 (no separate + generation worker; no colocation toggle) + Shared GPUs = cluster.gpus_per_node × cluster.num_nodes + + Worker groups, all colocated on the SAME GPUs: + teacher_0, teacher_1, …, teacher_{N-1}, student + + Per training step — serial time-slice on those GPUs: + teacher_0 : onload → forward → ship logits (CUDA IPC) → offload + teacher_1 : onload → forward → ship logits (CUDA IPC) → offload + ⋮ + teacher_{N-1} : onload → forward → ship logits (CUDA IPC) → offload + student : train → read each teacher's IPC buffer → release + + At any instant: exactly ONE teacher's params OR the student's resident, + plus the teachers' IPC logit buffers (node-local). +``` + +Future work will ease these requirements — we are actively working on +improving cross-tokenizer distillation support. + +### Step 1 — Build multi-token mappings + +Many student tokens (e.g., `"12"`) tokenize into multiple teacher tokens +(e.g., `"1"`, `"2"`). `minimal_projection_via_multitoken.py` walks the +student vocab, re-tokenizes each token with the teacher tokenizer, and adds +weighted entries to the projection. With `--enable-reverse-pass` it also +does the symmetric teacher → student walk. + +```bash +uv run python -m tools.x_token.minimal_projection_via_multitoken \ + --student-model "meta-llama/Llama-3.2-1B" \ + --teacher-model "Qwen/Qwen3-4B" \ + --top-k 32 \ + --enable-scale-trick \ + --enable-reverse-pass \ + --enable-special-token-mapping +``` + +Output: `cross_tokenizer_data/projection_map_Llama-3.2_to_Qwen3_multitoken_top_32_double_special.pt`. + +Pass `--num-examples 50` to print a sample of student→teacher mappings after +the matrix is built — useful for spot-checking that special tokens, numerals, +and punctuation map to sensible teacher tokens. + +### Step 2 — Reapply exact-token map + +Tokenizers built with a similar algorithm (for example, BPE) typically share +a sizable set of identical tokens — common punctuation, single ASCII +characters, and frequent subwords. `reapply_exact_map.py` pins those +overlapping tokens to 1-to-1 mappings with weight 1.0, overwriting whatever +Step 1 produced for them. + +```bash +uv run python -m tools.x_token.reapply_exact_map \ + --student-model "meta-llama/Llama-3.2-1B" \ + --teacher-model "Qwen/Qwen3-4B" \ + --initial-projection-path cross_tokenizer_data/projection_map_Llama-3.2_to_Qwen3_multitoken_top_32_double_special.pt +``` + +Output is written next to the input as `_exact_map_remapped.pt`. + +### Step 3 — Sort and trim to runtime `top_k` + +We observe the projection map is very sparse — each student token maps to at most 4–5 +teacher tokens. This step sorts each row by weight, trims to the chosen +runtime `top_k`, and stores the result as a sparse `[V_student, top_k]` +representation (per-row indices plus weights). That sparse format avoids materializing a computationally +expensive dense projection matrix of size `[student_vocab, teacher_vocab]` during distillation. + +```bash +uv run python -m tools.x_token.sort_and_cut_projection_matrix \ + --initial-projection-path cross_tokenizer_data/projection_map_Llama-3.2_to_Qwen3_multitoken_top_32_double_special_exact_map_remapped.pt \ + --top_k 4 \ + --output_path cross_tokenizer_data/projection_matrix_llama_qwen_top4.pt +``` + +### Step 4 — Launch x-token distillation + +The training entrypoint is `examples/run_xtoken_off_policy_distillation.py` with the +exemplar config at `examples/configs/xtoken_off_policy_distillation.yaml`. The exemplar +defaults to Llama-3.2-1B (student) ← Qwen3-4B (teacher) and the P-KL loss +mode. For data it points `data.train.data_files` at the ungated, CC-BY-4.0 +NVIDIA **Nemotron-Pretraining-Specialized-v1.1** corpus +(`Nemotron-Pretraining-Formal-Logic` subset) by dataset id + `data.train.subset`, +so the recipe runs out of the box with no auth or extra setup (and from the HF +cache under `HF_HUB_OFFLINE=1`). The `projection_matrix_path` below +points at the `cross_tokenizer_data/` directory that Steps 1–3 create, so run +those first (or the `build_projection_matrix.sh` wrapper). Override paths via +Hydra CLI: + +```bash +uv run python examples/run_xtoken_off_policy_distillation.py \ + --config examples/configs/xtoken_off_policy_distillation.yaml \ + teachers.0.aligner.projection_matrix_path=cross_tokenizer_data/projection_matrix_llama_qwen_top4.pt \ + cluster.gpus_per_node=8 \ + cluster.num_nodes=1 +``` + +The exemplar config keeps each teacher's `aligner.projection_matrix_path` as +`null`, so the projection matrix must always be supplied at the CLI (per +teacher, e.g. `teachers.0.aligner.projection_matrix_path=...`) — this keeps the +config reusable across (student, teacher) pairs. `data.train.data_files` and `data.train.subset` +already point at the default NVIDIA corpus described above; override them only +to train on your own `.arrow`/`.parquet`/`.json`/`.txt` corpus. Set both, since a +file path requires `subset` to be null: +`data.train.data_files=/path/glob/*.arrow data.train.subset=null`. + +#### Text and chat batches + +Both exemplar configs declare a typed top-level `collator` block: + +```yaml +collator: + mode: text + include_thinking_in_loss: false + native_thinking_alignment: false + kd_alignment_regions: null + num_packed_rows: 1 +``` + +For conversation data, set `collator.mode=chat` and use `chat_kd_processor` with +raw `messages` or `conversation` data. Each model applies its own chat template +to the full conversation, but training supervises **only the final assistant +turn**. Earlier messages remain context and do not contribute training targets. +Every chat row must end with an assistant message. CE, same-tokenizer KD, and +teacher scoring include the final assistant content and its retained end-of-turn +token; earlier turns, role headers, and padding are excluded. Cross-tokenizer KD +aligns the final assistant content and uses a separate EOT pair when both +tokenizations retain the terminator. This applies to both ordinary and native +chat alignment; raw-text training is unchanged. + +Each chat row can supply `message_loss_mask`, a list containing one binary +integer per message. Only assistant turns may have a value of `1`, and the final +assistant message must be selected (`1`). Earlier mask entries never enable +training on history: earlier messages are context only, even when their mask +entry is `1`. Omitting the mask selects the final assistant turn for training. +The OpenAI adapter prepends a zero when it inserts a system message. + +The configuration class supplies defaults. Packing multiple examples per row +remains unsupported; keep `num_packed_rows=1`. Text tokenization uses +`add_special_tokens=False`. Chat +rows exceeding either the student's or any teacher's context limit are rejected +instead of independently truncating their conversations. + +#### Native reasoning and tool conversations + +Native alignment supports the Nano/Qwen ChatML layouts with explicit reasoning, +answer, tool-call, and end-of-turn regions. Literal ChatML turn delimiters inside +source content or schemas are rejected as ambiguous; embedded tool/think markup +inside tool arguments is supported. Other layouts fail with an error; +the Llama text exemplars are not native-chat recipes. Enable it explicitly: + +```yaml +collator: + mode: chat + include_thinking_in_loss: true + native_thinking_alignment: true + kd_alignment_regions: [reasoning, close, answer, eot] + num_packed_rows: 1 +``` + +Each row can supply `tools`. Set `use_preserving_dataset: true` for JSONL tool +data to retain heterogeneous argument objects without Arrow adding null fields. + +The final assistant turn supervises reasoning text (when +`include_thinking_in_loss` is true), the closing ``, answers, tool payloads, +and EOT. Opening thinking scaffolds, formatting-only tokens, schemas, +user/system/tool-result text, earlier turns, and padding +are excluded. `kd_alignment_regions` further restricts cross-tokenizer KD without +changing CE or same-tokenizer KD masks. Prose and each tool call align separately; +permitted Boolean/null spelling differences retain CE but omit the unequal piece +from exact-text KD. Missing or transformed required content fails explicitly. +An empty or whitespace-only final assistant turn is rejected; tool-only final +turns are supported. Supervising a separate `reasoning_content` field requires native +alignment. Ordinary chat supports reasoning embedded in `content` only when the +chat template renders that content verbatim. For Qwen-style templates that extract +inline `...` into a separate reasoning block, including the stock +Qwen3 template and `qwen3_history.jinja`, set +`collator.native_thinking_alignment=true`. Ordinary chat rejects requests to +supervise a separate reasoning field without native alignment. + +Qwen3-4B's stock template drops historical reasoning when a later user turn +follows it. To retain that reasoning as context, apply the maintained override to +**each Qwen tokenizer** used by the collator (both `policy.tokenizer` and the +appropriate `teachers[i].tokenizer`): + +```yaml +tokenizer: + name: Qwen/Qwen3-4B + tokenizer_kwargs: + revision: 1cfa9a7208912126459214e8b04321603b3df60c + chat_template: examples/chat_templates/qwen3_history.jinja + chat_template_kwargs: + enable_thinking: true + preserve_thinking: true + truncate_history_thinking: false +``` + +The template accepts explicit `reasoning_content` and leading inline +`...` content. An explicit `truncate_history_thinking` wins: +`false` retains history and `true` uses stock history truncation. When absent, +`preserve_thinking=true` enables retention; otherwise stock behavior applies. +`enable_thinking` preserves Qwen3's generation-prompt behavior, including its +empty thinking block when false; it does not delete supplied reasoning. Training +renders with `add_generation_prompt=False`. Native training requires the final +turn's supervised reasoning to be retained and rejects a template that silently +drops it. Historical reasoning follows the template's context-retention policy +and does not contribute loss. + +The override adds explicit history-retention controls to the pinned Qwen3-4B +template while preserving its tool serialization and generation behavior. The shared offset matcher compares literal decoded strings; whitespace +differences are unequal. Native Unicode repair is separate and requires both +NFC-normalized decodings to equal the original source region. Ordinary chat uses +one canonical nonempty content key on both sides for the final assistant turn. +Content tokens that also absorb surrounding whitespace retain CE supervision; +their full decoded strings still determine exact-match correctness. + +The [`xtoken_off_policy_distillation.sh`](../../tests/functional/xtoken_off_policy_distillation.sh) +functional test covers raw text. The standalone +[`xtoken_native_chat.sh`](../../tests/functional/xtoken_native_chat.sh) functional +test covers native chat and is registered separately in the fast and full CI +suites. Its +[`chat configuration`](../../examples/configs/xtoken_native_chat.yaml) +uses a SmolLM2-135M-Instruct student and Qwen3-1.7B teacher with distinct +tokenizer vocabularies and token boundaries. Both use the history-preserving +ChatML override so the test retains reasoning, tool calls, and end-of-turn +supervision. The script reads the checked-in multi-turn arithmetic +[`JSONL fixture`](../../tests/functional/test_data/xtoken_native_chat.jsonl), builds +its projection matrix, and checks that three training steps produce finite +losses. It owns its experiment directory and logs, so it can run independently: + +```bash +bash tests/functional/xtoken_native_chat.sh +``` + +The chat scenario is functional coverage, not a nightly convergence benchmark. +DP alignment, native student-only SFT, sequence packing, and backend/loss changes +are outside this port. + +#### Loss-mode knobs + +`loss_fn` has two flags that pick between three behaviors: + +| `gold_loss` | `xtoken_loss` | Behavior | +|---|---|---| +| `false` | (inert) | **P-KL** — full-vocab teacher logits; the loss derives a microbatch-global top-k inside, projects the student into teacher vocab via the projection matrix, and chunk-averages KL on the top-k subset. CE term is added. | +| `true` | `false` | **Gold loss** — split the vocab into an *exact-token-mapped* common set (KL) and an *uncommon* tail (sorted L1). | +| `true` | `true` | **H-KL (gold + xtoken)** — same as gold, but relax the exact-map threshold to `>= 0.6` and allow multi-token projections to count as exact maps via a collision-replacement rule. | + +Other relevant fields: + +- `loss_fn.temperature` — softmax temperature applied symmetrically to student and teacher logits before KL. +- `loss_fn.vocab_topk` — microbatch-global top-k size for the P-KL path and for same-tokenizer teachers, whatever `gold_loss` is (the all-same-tokenizer `averaged_logits` path uses the full vocab). +- `loss_fn.uncommon_topk` — cap on the L1 uncommon-tail sort in the gold path (defaults to 8192). +- `loss_fn.reverse_kl` — compute `KL(student || teacher)` instead of `KL(teacher || student)`. + +### Results — 100-step multi-teacher run + +```bash +uv run python examples/run_xtoken_off_policy_distillation.py \ + --config examples/configs/xtoken_multiteacher_off_policy_distillation.yaml \ + teachers.0.aligner.projection_matrix_path=cross_tokenizer_data/projection_matrix_llama_phi-mini_top4.pt \ + cluster.gpus_per_node=8 \ + cluster.num_nodes=1 +``` + +This run distills a `meta-llama/Llama-3.2-1B` student from two teachers at +once: `microsoft/Phi-4-mini-instruct` (a cross-tokenizer teacher, projected +through its projection matrix) and `meta-llama/Llama-3.2-3B` (which shares the +student's tokenizer, so it contributes a direct top-K KL (`loss_fn.vocab_topk`) +with no projection). The per-teacher objectives are summed (`loss_fn.kd_loss_mode=sum`). +Config: global batch 96, micro-batch 1, sequence +length 2048, 100 steps, 2 nodes (8 GPUs each), on the default +Nemotron-Pretraining-Specialized-v1.1 / Formal-Logic corpus. The distillation +objective converges and the student tracks both teachers more closely over +training: + +train/loss, train/kl_loss, train/ce_loss, and train/accuracy over 100 multi-teacher distillation steps + +- **Loss** falls from ≈1.51 to ≈0.70. +- **KL loss** (summed over both teachers) falls from ≈4.89 to ≈1.97. Almost + all of it comes from the cross-tokenizer Phi-4-mini teacher (≈4.75 → ≈1.76); + the same-tokenizer Llama-3.2-3B teacher's direct KL is already small and + stays there (≈0.14 → ≈0.21). +- **CE loss** falls from ≈0.75 to ≈0.35. +- **Top-1 accuracy** rises from ≈0.82 to ≈0.91. + +#### Downstream evaluation + +Benchmark scores of the distilled student against the undistilled +`meta-llama/Llama-3.2-1B` base model: + +| Task | Base Llama-3.2-1B | Distilled | +|---|---|---| +| MMLU | 32.05 | 40.24 | +| GSM8K | 5.69 | 5.76 | + +#### Throughput and memory + +Measured on the same run (per training step, micro-batch 1, sequence length +2048, with the two teachers forwarded serially): + +| Metric | Value | +|---|---| +| Mean step time | 6.65 s (min 6.22 s) | +| Teacher forward (both teachers) | 4.94 s mean — the dominant per-step cost | +| Training throughput | ≈29.5k valid tokens/s (196,512 tokens/step ÷ mean step time) | + +Each teacher's full-vocab logits stay on-node: the producer publishes a +rank-level `[B_r, T_t, V_t]` bf16 tray and hands the student a CUDA IPC handle +to it, so teacher logits never cross the network even with two teachers in +play. + +## CSD (Common-Support Distillation) + +CSD compares the two models only on events that both tokenizers can express. A **REST** category holds the remaining probability mass, so each comparison stays a proper distribution. + +The shared events come from two static **sub-token tables** (support tables) built once per (student, teacher) tokenizer pair: a student-to-teacher table and a teacher-to-student table. Each table maps a source token to the complete chain of destination tokens that spells the same text. + +This section distills a `Qwen/Qwen3-4B` teacher into a `meta-llama/Llama-3.2-1B` student and explains how to: + +1. Build and validate the sub-token tables for a (student, teacher) tokenizer pair. +2. Launch CSD distillation with the tables. + +### How CSD works + +A full run has two phases. The three preparation steps run once per (student, teacher) pair on CPU — they load tokenizers only, never model weights — and produce two `.pt` files. The final step is the distillation training loop. + +``` + ┌──────────────────────────────────────────────┐ + │ Offline sub-token table preparation (CPU) │ + │ │ + │ ┌────────────────────────────────────┐ │ + (student, teacher) │ │ 1. build_subtoks_table.py │ │ + tokenizers ────▶│ │ — student → teacher chains │ │ + │ └─────────────────┬──────────────────┘ │ + │ │ │ + │ ┌─────────────────▼──────────────────┐ │ + │ │ 2. build_subtoks_table.py │ │ + │ │ — teacher → student chains │ │ + │ └─────────────────┬──────────────────┘ │ + │ │ │ + │ ┌─────────────────▼──────────────────┐ │ + │ │ 3. validate_subtoks_tables.py │ │ + │ │ — check both against tokenizers │ │ + │ └─────────────────┬──────────────────┘ │ + └────────────────────│─────────────────────────┘ + │ + ▼ student_to_teacher.pt + teacher_to_student.pt + ┌────────────────────────────────────────────────────┐ + │ 4. examples/run_csd.py │ + │ — align student & teacher tokens, then │ + │ teacher forward + student forward, │ + │ then CSD loss + student CE │ + └────────────────────────────────────────────────────┘ +``` + +In every training step, the collator tokenizes each text with both tokenizers and aligns their character spans into chunks. A chunk covers `M` student tokens and `N` teacher tokens that spell the same text. The teacher runs a frozen forward pass and sends its logits to the student over CUDA IPC; the student then computes the CSD loss on each chunk together with its own next-token cross-entropy (CE). Only the student is updated. + +```mermaid +flowchart TD + text["Raw text"] --> tokens["Tokenize with the student and teacher tokenizers"] + tokens --> align["Align character spans into chunks"] + align --> forward["Frozen teacher forward and student forward"] + forward --> kind{"Chunk type"} + tables["Sub-token tables"] -.-> common + tables -.-> prefix + kind -->|"1:1"| common["Common tokens"] + kind -->|"1:many or many:1"| mode{"csd_realized_only_mismatch"} + kind -->|"many:many"| realized["Realized token chains"] + mode -->|"false (default)"| prefix["Prefix-support alternatives"] + mode -->|"true"| realized + common --> rest["Add the REST bucket"] + prefix --> rest + realized --> rest + rest --> divergence["Per-chunk divergence (JSD by default)"] + divergence --> kd["Weight mismatch chunks by csd_mismatch_loss_beta
Normalize by the global chunk count"] + forward --> ce["Student next-token CE"] + kd --> total["Total loss"] + ce --> total + total --> update["Update the student only"] +``` + +Each chunk type compares a different set of events: + +| Chunk (student `M` × teacher `N` tokens) | Compared events (default) | With `csd_realized_only_mismatch: true` | +|---|---|---| +| 1:1 | Common tokens — vocabulary entries that map one-to-one between the tokenizers — plus REST | Same | +| 1:many | Student tokens whose teacher chain has length `N` and shares the realized teacher prefix, plus REST. The teacher scores an alternative as the realized-prefix chain log-probability plus the alternative's final-token log-probability | Realized chains plus REST | +| many:1 | Mirror of 1:many, using the teacher-to-student table | Realized chains plus REST | +| many:many | Realized student chain against realized teacher chain, plus REST | Same | + +With sparse teacher transport, the teacher sends its top-K logits per row together with the exact full-vocabulary log-normalizer, so REST holds the true remaining probability mass: top-K is not a renormalization over K tokens. The realized label always stays in the compared support, and the alternatives are limited to the teacher's top-K. + +The KD term for teacher `i` is + +``` +KD_i = T² × (Σ 1:1 chunk losses + csd_mismatch_loss_beta × Σ mismatch chunk losses) / (valid chunks in the global batch) +``` + +and the total loss is `kl_loss_weight × KD + ce_loss_scale × CE`. With `dynamic_loss_scaling: true` (the exemplar), the KD term is instead rescaled to the detached CE magnitude, `sg(|CE| / |KD|) × KD + CE`, and both weights are ignored. A teacher configured without tables shares the student's tokenizer: it skips alignment and contributes a top-`vocab_topk` KL on the shared vocabulary. + +### CSD backend and scope + +- **Automodel or Megatron, per model.** The student and each teacher choose their backend independently through `automodel_cfg.enabled` or `megatron_cfg.enabled`. The exemplar runs both models on Megatron (TP1/CP1/PP1) on one node with eight GPUs and offloads the student after each step. +- **Teacher logits travel via CUDA IPC**, so the student and teachers are colocated on the same node. On more than one node, every teacher matches the student's DP size and TP × CP × PP group size, and each group fits inside one node. +- **One colocated GPU pool, serial execution.** As with X-Token, the teachers and the student time-share one GPU pool within each step; see the [X-Token backend notes](#backend-and-scope). + +Each cross-tokenizer teacher sends its logits in one of two forms: + +- **Sparse** (`loss_fn.teacher_topk_ipc_k > 0`, default `128`): for each teacher row, the top-K logits, the exact log-normalizer, and the realized teacher labels the loss needs. +- **Dense** (`teacher_topk_ipc_k: 0`): the full-vocabulary logits. + +Same-tokenizer teachers always send dense logits. Sparse and dense transport optimize different objectives: with sparse transport, the KD compares each teacher row's top-K tokens plus REST instead of the full vocabulary. + +A Megatron student evaluates cross-tokenizer KD **owner-local** when the run allows it: each context-parallel (CP) rank keeps its own share of the student logits, evaluates the chunks whose final student predictor row it holds, and reads only the teacher rows those chunks need. Otherwise the loss is evaluated **replicated**: every CP rank rebuilds its contiguous sequence window of the student and teacher logits. + +| Student | Cross-tokenizer teacher | `teacher_topk_ipc_k` | Teacher transport | KD evaluation | +|---|---|---|---|---| +| Megatron, owner-local eligible | Megatron | `> 0` | Sparse rows of each teacher CP shard | Owner-local | +| Megatron, owner-local eligible | Megatron or Automodel | `0` | Dense | Owner-local | +| Megatron, not eligible | Megatron or Automodel | `0` | Dense | Replicated | +| Automodel | Automodel | `> 0` | Sparse | Replicated | +| Automodel | Megatron or Automodel | `0` | Dense | Replicated | + +Setup rejects the remaining combinations and prints which evaluation a run uses and why; set `teacher_topk_ipc_k: 0` for those runs. Owner-local evaluation requires: + +- a Megatron student with `pipeline_model_parallel_size: 1`, no sequence packing, and no dynamic batching; +- `kd_loss_mode: sum` with static teacher weights (`sum_weights_metric: null`); +- `csd_pure_alm: false` and a `csd_loss_fn` other than `bce`; `csd_last_pos_loss_fn: bce` only with `csd_realized_only_mismatch: true`; +- with sparse transport, `teacher_topk_ipc_keep_realized: true` (or realized-only mismatch), Megatron cross-tokenizer teachers, and a teacher `make_sequence_length_divisible_by` that is a multiple of `2 × CP` (and of TP when the teacher uses sequence parallelism). + +Under owner-local evaluation the student logits are never gathered across CP; only per-row scalars and the selected token values cross ranks. `csd_selected_row_chunk_size` and `csd_support_chunk_batch_size` bound the temporary memory of this evaluation without changing the loss. + +### Step 1 — Build the student-to-teacher table + +`build_subtoks_table.py` builds one direction per call. The forward table has student token IDs as rows and teacher token IDs as entries: + +```bash +uv run python -m tools.csd.build_subtoks_table \ + --source-tokenizer meta-llama/Llama-3.2-1B \ + --destination-tokenizer Qwen/Qwen3-4B \ + --max-subtokens 8 \ + --output cross_tokenizer_data/csd/student_to_teacher.pt +``` + +Each `.pt` file stores, for every source token, the complete destination chain: + +| Key | Meaning | +|---|---| +| `subtoks` | Integer tensor `[V_src, N_max]` of destination token IDs; unused entries are `-1` | +| `lengths` | Integer vector of chain lengths; `0` marks an unmapped source token | +| `V_src`, `V_dst`, `N_max` | Vocabulary sizes (`len(tokenizer)`, including added tokens) and the chain-length limit | +| `src`, `dst` | Source and destination tokenizer identifiers | +| `source_fingerprint`, `destination_fingerprint` | SHA-256 fingerprints of each tokenizer's backend and special-token roles | + +For example, if source token `42` is `playing` and the destination tokenizer spells it as `play` (`7`) + `ing` (`9`), then with `--max-subtokens 8`: + +```text +lengths[42] = 2 +subtoks[42] = [7, 9, -1, -1, -1, -1, -1, -1] +``` + +These IDs are illustrative. A token shared by both vocabularies has a length-one chain, and the length-one chains of the forward table form the common-token map when `loss_fn.common_indices_from_subtoks: true` (the exemplar). The tables describe vocabulary events only; training still tokenizes each text with both tokenizers and aligns the actual spans. + +The builder creates the output directory. Use `--source-revision` and `--destination-revision` to pin tokenizer revisions, and `--overwrite` to replace an existing file. + +### Step 2 — Build the teacher-to-student table + +Reverse the tokenizer order to build the reverse table, which has teacher token IDs as rows and student token IDs as entries: + +```bash +uv run python -m tools.csd.build_subtoks_table \ + --source-tokenizer Qwen/Qwen3-4B \ + --destination-tokenizer meta-llama/Llama-3.2-1B \ + --max-subtokens 8 \ + --output cross_tokenizer_data/csd/teacher_to_student.pt +``` + +The two directions are separate decompositions: a forward chain is not necessarily inverted by a single reverse row. The 1:many chunks use the forward table and the many:1 chunks use the reverse table. + +The builder maps each source token as follows: + +- **Special tokens** map by **role** — BOS, EOS, PAD, UNK, SEP, CLS and MASK — even when their spellings differ. When roles do not decide, identical special-token spellings map only if their roles do not conflict. Any other special token stays unmapped; it never decomposes into ordinary text. +- **Ordinary tokens** first use an exact surface match. Otherwise, the token is decoded and re-encoded with the destination tokenizer; for byte-level BPE pairs, the result must keep the original byte surface, and BPE on that surface preserves incomplete UTF-8 fragments. Chains that contain a destination special token or a U+FFFD replacement character are rejected. +- **Long chains** — longer than `--max-subtokens` (default `8`) — stay unmapped and are never truncated. Raise the limit to keep longer decompositions; tables grow accordingly. + +The builder and validator support fast tokenizers with a **BPE model and a ByteLevel decoder**, including the example pair. Metaspace (SentencePiece), WordPiece and other decoder families are rejected, because decoding one token in isolation can change whitespace or word markers. + +### Step 3 — Validate both tables + +Check both tables against the exact student and teacher tokenizers before allocating GPUs: + +```bash +uv run python -m tools.csd.validate_subtoks_tables \ + --student-tokenizer meta-llama/Llama-3.2-1B \ + --teacher-tokenizer Qwen/Qwen3-4B \ + --student-to-teacher cross_tokenizer_data/csd/student_to_teacher.pt \ + --teacher-to-student cross_tokenizer_data/csd/teacher_to_student.pt +``` + +Pass `--student-revision` and `--teacher-revision` when the tables were built with pinned revisions. The validator checks tokenizer fingerprints, tensor shapes and dtypes, destination ID ranges, padding, special-token roles, and every mapped event, and exits with an error on an inconsistent table. Training does not re-check the tables against the runtime tokenizers, so validate every table pair and rebuild both directions whenever either tokenizer changes. Tables are loaded with `torch.load(weights_only=False)`; use only trusted files. + +To inspect the mapping coverage of a pair: + +```bash +uv run python - <<'PY' +import torch + +for name in ("student_to_teacher", "teacher_to_student"): + table = torch.load(f"cross_tokenizer_data/csd/{name}.pt", weights_only=True) + lengths = table["lengths"] + print( + f"{name}: shape={tuple(table['subtoks'].shape)} " + f"mapped={int((lengths > 0).sum())} unmapped={int((lengths == 0).sum())} " + f"direct={int((lengths == 1).sum())} chains={int((lengths > 1).sum())}" + ) +PY +``` + +Some unmapped rows are expected, for example a BOS token with no counterpart role in the other tokenizer. + +### Step 4 — Launch CSD distillation + +The training entrypoint is `examples/run_csd.py`, which loads the exemplar `examples/configs/csd.yaml` by default. The exemplar distills `Qwen/Qwen3-4B` into `meta-llama/Llama-3.2-1B` on the ungated, CC-BY-4.0 **Nemotron-Pretraining-Specialized-v1.1** corpus (`Nemotron-Pretraining-Formal-Logic` subset), and its table paths already point at the files written in Steps 1 and 2: + +```bash +uv run python examples/run_csd.py \ + cluster.gpus_per_node=8 \ + cluster.num_nodes=1 +``` + +The table settings in the exemplar are: + +```yaml +algorithm: csd +loss_fn: + common_indices_from_subtoks: true +teachers: + - model_name: Qwen/Qwen3-4B + tokenizer: + name: Qwen/Qwen3-4B + aligner: + projection_matrix_path: null + student_to_teacher_table_path: cross_tokenizer_data/csd/student_to_teacher.pt + teacher_to_student_table_path: cross_tokenizer_data/csd/teacher_to_student.pt +``` + +For tables stored elsewhere, override `teachers.0.aligner.student_to_teacher_table_path` and `teachers.0.aligner.teacher_to_student_table_path`. The table files must be readable at those paths on every training worker. Models and tables stay paired: to change the student, override both `policy.model_name` and `policy.tokenizer.name`; to change a teacher, override both `teachers.0.model_name` and `teachers.0.tokenizer.name`; and in either case rebuild and revalidate both tables. + +To train on your own corpus, set `data.train.data_files` to a file, glob, or dataset identifier and `data.train.text_key` to its text column; for a file or glob, also set `data.train.subset=null`. Supported local formats include `.arrow`, `.parquet`, `.json`, `.jsonl` and `.txt`. In text mode (`collator.mode: text`), a row that does not fit every model's sequence limit is clipped to the longest text prefix that fits all of them, so the student and every teacher still see the same text. For conversational data, see [Text and chat batches](#text-and-chat-batches); CSD chat mode supervises only the final assistant turn: CE and KD both use that turn's content and end-of-turn token, and template markers such as the role header are not supervised. + +Checkpoints are saved in the student backend's format; convert a Megatron checkpoint to Hugging Face format with the command in [Evaluation](../about/evaluation.md). + +#### Loss knobs + +| Field | Purpose | +|---|---| +| `csd_loss_fn` | Divergence on 1:1 chunks: `kl`, `jsd` (default) or `bce` | +| `csd_last_pos_loss_fn` | Divergence on mismatch chunks; `null` inherits `csd_loss_fn` | +| `csd_jsd_beta` | JSD mixture coefficient: `0` is forward KL, `1` is reverse KL, `0.5` (default) is symmetric | +| `csd_mismatch_loss_beta` | Weight of each 1:many, many:1 and many:many chunk relative to a 1:1 chunk (default `5.0`) | +| `csd_realized_only_mismatch` | Score 1:many and many:1 chunks by their realized chains against REST instead of prefix-support alternatives | +| `csd_pure_alm` | Score every chunk, including 1:1, by its realized token or chain against REST | +| `csd_noise_filter_topk` | Skip a chunk when a realized teacher label falls outside the teacher's top-k at its predictor (`0` disables) | +| `kl_chunk_shift` | Predict each non-initial chunk from the position one step earlier (default `true`) | +| `teacher_topk_ipc_k` | Teacher logits sent per row: top-K with exact normalization (default `128`), or `0` for dense full-vocabulary logits | +| `temperature`, `reverse_kl` | Softmax temperature for both models; reverse KL for the `kl` divergence | +| `dynamic_loss_scaling`, `kl_loss_weight`, `ce_loss_scale` | KD/CE combination; the weights apply only when `dynamic_loss_scaling: false` | +| `vocab_topk` | Top-K vocabulary for same-tokenizer teachers | + +#### Multiple teachers + +Cross-tokenizer and same-tokenizer teachers can be mixed in the `teachers` list; `kd_loss_mode` combines their KD terms: + +| `kd_loss_mode` | Behavior | +|---|---| +| `sum` | `Σ_i w_i × KD_i` with static weights `teachers[i].weight`, or with `sum_weights_metric` (`ce`, `entropy`, `max_prob`) for weights `softmax(alpha × score)` | +| `averaged_logits` | One full-vocabulary KL against the weighted average of teacher logits when every teacher is same-tokenizer; otherwise a static weighted sum | +| `select_teacher` | Per microbatch, only the teacher with the lowest next-token CE on its own tokens | + +`select_teacher` and `sum_weights_metric` need dense transport (`teacher_topk_ipc_k: 0`). + +`examples/configs/csd_multiteacher.yaml` pairs a cross-tokenizer `microsoft/Phi-4-mini-instruct` teacher with a same-tokenizer `meta-llama/Llama-3.2-3B` teacher. Build and validate a Llama-3.2-1B / Phi-4-mini-instruct table pair (Steps 1–3) at their own paths and pass them as overrides: + +```bash +uv run python examples/run_csd.py \ + --config examples/configs/csd_multiteacher.yaml \ + teachers.0.aligner.student_to_teacher_table_path=cross_tokenizer_data/csd/phi4mini/student_to_teacher.pt \ + teachers.0.aligner.teacher_to_student_table_path=cross_tokenizer_data/csd/phi4mini/teacher_to_student.pt +``` + +#### Optional: lockstep sequence packing + +CSD can pack sequences when the student and every teacher pack the **same samples into the same bins**, each with its own tokenization. Enable it on every side: + +```bash +uv run python examples/run_csd.py \ + loss_fn.teacher_topk_ipc_k=0 \ + policy.sequence_packing.enabled=true \ + policy.sequence_packing.algorithm=lockstep_first_fit_decreasing \ + teachers.0.sequence_packing.enabled=true \ + teachers.0.sequence_packing.algorithm=lockstep_first_fit_decreasing +``` + +Lockstep packing uses dense transport, `kd_loss_mode: sum` with static weights, micro-batch size 1, and equal DP sizes on every side; setup validation rejects other combinations before any worker starts. The [lockstep packing design](../design-docs/csd-lockstep-sequence-packing.md) lists the full support matrix. + + + +## Where files live + +| Stage | Tool | Default output | +|---|---|---| +| Build multi-token | `tools/x_token/minimal_projection_via_multitoken.py` | `/projection_map__to__multitoken_top__double[_special].pt` | +| Reapply exact map | `tools/x_token/reapply_exact_map.py` | `_exact_map_remapped.pt` | +| Sort and trim | `tools/x_token/sort_and_cut_projection_matrix.py` | `/_top__sorted[_preservelast].pt` (or `--output_path`) | +| Train (X-Token) | `examples/run_xtoken_off_policy_distillation.py` | per the run's `logger.log_dir` and `checkpointing.checkpoint_dir` | +| Build student-to-teacher table | `tools/csd/build_subtoks_table.py` | `--output` (exemplar: `cross_tokenizer_data/csd/student_to_teacher.pt`) | +| Build teacher-to-student table | `tools/csd/build_subtoks_table.py` | `--output` (exemplar: `cross_tokenizer_data/csd/teacher_to_student.pt`) | +| Validate | `tools/csd/validate_subtoks_tables.py` | None; exits non-zero on an inconsistent table | +| Train (CSD) | `examples/run_csd.py` | Per the run's `logger.log_dir` and `checkpointing.checkpoint_dir` | + +The nightly CSD suites (`tests/test_suites/llm/csd-*.sh`) build both tables in the experiment directory before training and pass their paths as overrides; they are a working reference for scripted preparation. + +## Related + +Listed in the order they run during a training step: + +- Config exemplars: [`examples/configs/xtoken_off_policy_distillation.yaml`](../../examples/configs/xtoken_off_policy_distillation.yaml), [`examples/configs/csd.yaml`](../../examples/configs/csd.yaml) and [`examples/configs/csd_multiteacher.yaml`](../../examples/configs/csd_multiteacher.yaml) +- Trainer module: `nemo_rl/algorithms/xtoken_off_policy_distillation.py` +- KD data processor: `nemo_rl/data/processors.py::kd_data_processor` +- Cross-tokenizer collators: `nemo_rl/data/cross_tokenizer/cross_tokenizer_collate.py::CrossTokenizerCollator` and its CSD subclass `nemo_rl/data/cross_tokenizer/csd_collate.py::CSDCrossTokenizerCollator` +- Token alignment: `nemo_rl/algorithms/x_token/token_aligner.py::TokenAligner` +- Loss implementations: `nemo_rl/algorithms/loss/loss_functions.py::CrossTokenizerDistillationLossFn` (X-Token) and `nemo_rl/algorithms/loss/csd.py::CSDLossFn` (CSD) +- CSD lockstep packing: [design doc](../design-docs/csd-lockstep-sequence-packing.md) and `nemo_rl/algorithms/loss/csd_packing.py` diff --git a/docs/guides/xtoken-off-policy-distillation.md b/docs/guides/xtoken-off-policy-distillation.md deleted file mode 100644 index 571f57db2ce..00000000000 --- a/docs/guides/xtoken-off-policy-distillation.md +++ /dev/null @@ -1,453 +0,0 @@ -# Cross-Tokenizer (X-Token) Off-Policy Distillation - -NeMo RL supports off-policy distillation between a student and a teacher that -**do not share a tokenizer** — for example, distilling a Qwen3-4B teacher into -a Llama-3.2-1B student. Cross-tokenizer ("x-token") distillation handles the -vocabulary mismatch by routing student logits through a precomputed -**projection matrix** that maps each student token to the teacher tokens it -most plausibly corresponds to, projecting the student into the teacher's -vocab space so the two distributions can be compared. - -Tokenizer vocabularies overlap only partially, which is what makes the -projection necessary. The table below reports the pairwise overlap -(intersection divided by the smaller vocabulary, on canonical token forms) -across several model tokenizers; the off-diagonal entries sit well below -`1.0` — that shared-vocabulary gap is exactly what the projection matrix -bridges. - -| Model | Mistral-NeMo-Minitron-8B | Qwen3-8B-Base | Llama-3.2-1B | gemma-3-4b-it | OLMo2-8B-SuperBPE-t160k | gpt-oss-20b | -|---|---:|---:|---:|---:|---:|---:| -| Mistral-NeMo-Minitron-8B | 1.0000 | 0.5119 | 0.5525 | 0.7430 | 0.4215 | 0.7591 | -| Qwen3-8B-Base | 0.5119 | 1.0000 | 0.8481 | 0.6320 | 0.4103 | 0.6462 | -| Llama-3.2-1B | 0.5525 | 0.8481 | 1.0000 | 0.6739 | 0.4980 | 0.7977 | -| gemma-3-4b-it | 0.7430 | 0.6320 | 0.6739 | 1.0000 | 0.4918 | 0.6545 | -| OLMo2-8B-SuperBPE-t160k | 0.4215 | 0.4103 | 0.4980 | 0.4918 | 1.0000 | 0.3657 | -| gpt-oss-20b | 0.7591 | 0.6462 | 0.7977 | 0.6545 | 0.3657 | 1.0000 | - -This guide explains how to: - -1. Create the projection matrix from a (student, teacher) tokenizer pair. -2. Launch a distillation with the projection matrix. - -## How it works - -A full run has two phases. The three prep steps are *offline data prep* — -small CLI tools you run once per (student, teacher) pair — and the result is a -single `.pt` file. The final step is the actual distillation training loop. - -``` - ┌──────────────────────────────────────────────┐ - │ Offline projection-matrix preparation │ - │ │ - │ ┌────────────────────────────────────┐ │ - (student, teacher) │ │ 1. minimal_projection_via_ │ │ - tokenizers ────▶│ │ multitoken.py │ │ - │ │ — multi-token mappings │ │ - │ └─────────────────┬──────────────────┘ │ - │ │ │ - │ ┌─────────────────▼──────────────────┐ │ - │ │ 2. reapply_exact_map.py │ │ - │ │ — pin exact 1-to-1 matches │ │ - │ └─────────────────┬──────────────────┘ │ - │ │ │ - │ ┌─────────────────▼──────────────────┐ │ - │ │ 3. sort_and_cut_projection_matrix │ │ - │ │ .py — trim to runtime top_k │ │ - │ └─────────────────┬──────────────────┘ │ - └────────────────────│─────────────────────────┘ - │ - ▼ projection_matrix.pt - ┌────────────────────────────────────────────────────┐ - │ 4. examples/ │ - │ run_xtoken_off_policy_distillation.py │ - │ — align student & teacher tokens, then │ - │ teacher forward + student forward, │ - │ then x-token KD loss │ - └────────────────────────────────────────────────────┘ -``` - -The projection matrix is a sparse `[V_student, top_k]` tensor that the -training-time loss multiplies against the student logits to project them into -the teacher's vocab space. - -Projection-matrix weights mapping student tokens (rows) to teacher tokens (columns) - -Each row of the matrix holds the weights `W_{s,t}` that distribute a student -token `s ∈ V_S` over the teacher tokens `t ∈ V_T` it corresponds to. Tokens -shared by both vocabularies map 1-to-1 (e.g., `_the`, `_cat`, `_run`), while a -student token that the teacher splits into pieces spreads its weight across -those pieces (e.g., `201` → `2`, `0`, `1`). Rows are trimmed to the runtime -`top_k` in Step 3, so low-weight tail entries are dropped (hatched cell). - -## Quickstart — single command - -For the typical case, `tools/x_token/build_projection_matrix.sh` chains -the prep steps with auto-derived intermediate paths: - -```bash -./tools/x_token/build_projection_matrix.sh \ - --student-model meta-llama/Llama-3.2-1B \ - --teacher-model Qwen/Qwen3-4B \ - --runtime-top-k 4 -``` - -The wrapper writes the final matrix to -`cross_tokenizer_data/projection_matrix___top.pt` -(override with `--final-output`). Further tweaks to step 1 defaults can be configured using `--no-{scale-trick,reverse-pass,special-token-mapping}`. Run `./tools/x_token/build_projection_matrix.sh ---help` for the full list of options. - -The per-step recipes below are for advanced customization (non-default -weight thresholds, hand-picked intermediate filenames, etc.). - -## Backend and scope - -- **Automodel only.** Set `policy.automodel_cfg.enabled=true`. The Megatron policy - worker is not wired for cross-tokenizer distillation. -- **Teacher logits travel via CUDA IPC**, so student and teacher policies must - be colocated on the same node. No remote-Ray transport for x-token logits. -- **One colocated GPU pool, serial execution.** All teachers and the student - share a single `RayVirtualCluster` - (`max_colocated_worker_groups = len(teachers) + 1`). There is no separate - generation worker (off-policy distillation trains on a fixed dataset) and no - colocation toggle. Each step time-slices the shared GPUs: one teacher at a - time is onloaded, runs its forward pass, ships its logits to the student over - CUDA IPC, then offloads; the student trains last, reading and releasing each - teacher's IPC buffer. GPU memory therefore holds exactly one teacher's params - *or* the student's at any instant, plus the teachers' resident IPC logit - buffers. - -``` - RayVirtualCluster — one GPU pool, shared by every worker group - max_colocated_worker_groups = len(teachers) + 1 (no separate - generation worker; no colocation toggle) - Shared GPUs = cluster.gpus_per_node × cluster.num_nodes - - Worker groups, all colocated on the SAME GPUs: - teacher_0, teacher_1, …, teacher_{N-1}, student - - Per training step — serial time-slice on those GPUs: - teacher_0 : onload → forward → ship logits (CUDA IPC) → offload - teacher_1 : onload → forward → ship logits (CUDA IPC) → offload - ⋮ - teacher_{N-1} : onload → forward → ship logits (CUDA IPC) → offload - student : train → read each teacher's IPC buffer → release - - At any instant: exactly ONE teacher's params OR the student's resident, - plus the teachers' IPC logit buffers (node-local). -``` - -Future work will ease these requirements — we are actively working on -improving cross-tokenizer distillation support. - -## Step 1 — Build multi-token mappings - -Many student tokens (e.g., `"12"`) tokenize into multiple teacher tokens -(e.g., `"1"`, `"2"`). `minimal_projection_via_multitoken.py` walks the -student vocab, re-tokenizes each token with the teacher tokenizer, and adds -weighted entries to the projection. With `--enable-reverse-pass` it also -does the symmetric teacher → student walk. - -```bash -uv run python -m tools.x_token.minimal_projection_via_multitoken \ - --student-model "meta-llama/Llama-3.2-1B" \ - --teacher-model "Qwen/Qwen3-4B" \ - --top-k 32 \ - --enable-scale-trick \ - --enable-reverse-pass \ - --enable-special-token-mapping -``` - -Output: `cross_tokenizer_data/projection_map_Llama-3.2_to_Qwen3_multitoken_top_32_double_special.pt`. - -Pass `--num-examples 50` to print a sample of student→teacher mappings after -the matrix is built — useful for spot-checking that special tokens, numerals, -and punctuation map to sensible teacher tokens. - -## Step 2 — Reapply exact-token map - -Tokenizers built with a similar algorithm (for example, BPE) typically share -a sizable set of identical tokens — common punctuation, single ASCII -characters, and frequent subwords. `reapply_exact_map.py` pins those -overlapping tokens to 1-to-1 mappings with weight 1.0, overwriting whatever -Step 1 produced for them. - -```bash -uv run python -m tools.x_token.reapply_exact_map \ - --student-model "meta-llama/Llama-3.2-1B" \ - --teacher-model "Qwen/Qwen3-4B" \ - --initial-projection-path cross_tokenizer_data/projection_map_Llama-3.2_to_Qwen3_multitoken_top_32_double_special.pt -``` - -Output is written next to the input as `_exact_map_remapped.pt`. - -## Step 3 — Sort and trim to runtime `top_k` - -We observe the projection map is very sparse — each student token maps to at most 4–5 -teacher tokens. This step sorts each row by weight, trims to the chosen -runtime `top_k`, and stores the result as a sparse `[V_student, top_k]` -representation (per-row indices plus weights). That sparse format avoids materializing a computationally -expensive dense projection matrix of size `[student_vocab, teacher_vocab]` during distillation. - -```bash -uv run python -m tools.x_token.sort_and_cut_projection_matrix \ - --initial-projection-path cross_tokenizer_data/projection_map_Llama-3.2_to_Qwen3_multitoken_top_32_double_special_exact_map_remapped.pt \ - --top_k 4 \ - --output_path cross_tokenizer_data/projection_matrix_llama_qwen_top4.pt -``` - -## Step 4 — Launch x-token distillation - -The training entrypoint is `examples/run_xtoken_off_policy_distillation.py` with the -exemplar config at `examples/configs/xtoken_off_policy_distillation.yaml`. The exemplar -defaults to Llama-3.2-1B (student) ← Qwen3-4B (teacher) and the P-KL loss -mode. For data it points `data.train.data_files` at the ungated, CC-BY-4.0 -NVIDIA **Nemotron-Pretraining-Specialized-v1.1** corpus -(`Nemotron-Pretraining-Formal-Logic` subset) by dataset id + `data.train.subset`, -so the recipe runs out of the box with no auth or extra setup (and from the HF -cache under `HF_HUB_OFFLINE=1`). The `projection_matrix_path` below -points at the `cross_tokenizer_data/` directory that Steps 1–3 create, so run -those first (or the `build_projection_matrix.sh` wrapper). Override paths via -Hydra CLI: - -```bash -uv run python examples/run_xtoken_off_policy_distillation.py \ - --config examples/configs/xtoken_off_policy_distillation.yaml \ - teachers.0.aligner.projection_matrix_path=cross_tokenizer_data/projection_matrix_llama_qwen_top4.pt \ - cluster.gpus_per_node=8 \ - cluster.num_nodes=1 -``` - -The exemplar config keeps each teacher's `aligner.projection_matrix_path` as -`null`, so the projection matrix must always be supplied at the CLI (per -teacher, e.g. `teachers.0.aligner.projection_matrix_path=...`) — this keeps the -config reusable across (student, teacher) pairs. `data.train.data_files` and `data.train.subset` -already point at the default NVIDIA corpus described above; override them only -to train on your own `.arrow`/`.parquet`/`.json`/`.txt` corpus. Set both, since a -file path requires `subset` to be null: -`data.train.data_files=/path/glob/*.arrow data.train.subset=null`. - -### Text and chat batches - -Both exemplar configs declare a typed top-level `collator` block: - -```yaml -collator: - mode: text - include_thinking_in_loss: false - native_thinking_alignment: false - kd_alignment_regions: null - num_packed_rows: 1 -``` - -For conversation data, set `collator.mode=chat` and use `chat_kd_processor` with -raw `messages` or `conversation` data. Each model applies its own chat template -to the full conversation, but training supervises **only the final assistant -turn**. Earlier messages remain context and do not contribute training targets. -Every chat row must end with an assistant message. CE, same-tokenizer KD, and -teacher scoring include the final assistant content and its retained end-of-turn -token; earlier turns, role headers, and padding are excluded. Cross-tokenizer KD -aligns the final assistant content and uses a separate EOT pair when both -tokenizations retain the terminator. This applies to both ordinary and native -chat alignment; raw-text training is unchanged. - -Each chat row can supply `message_loss_mask`, a list containing one binary -integer per message. Only assistant turns may have a value of `1`, and the final -assistant message must be selected (`1`). Earlier mask entries never enable -training on history: earlier messages are context only, even when their mask -entry is `1`. Omitting the mask selects the final assistant turn for training. -The OpenAI adapter prepends a zero when it inserts a system message. - -The configuration class supplies defaults. Packing multiple examples per row -remains unsupported; keep `num_packed_rows=1`. Text tokenization uses -`add_special_tokens=False`. Chat -rows exceeding either the student's or any teacher's context limit are rejected -instead of independently truncating their conversations. - -### Native reasoning and tool conversations - -Native alignment supports the Nano/Qwen ChatML layouts with explicit reasoning, -answer, tool-call, and end-of-turn regions. Literal ChatML turn delimiters inside -source content or schemas are rejected as ambiguous; embedded tool/think markup -inside tool arguments is supported. Other layouts fail with an error; -the Llama text exemplars are not native-chat recipes. Enable it explicitly: - -```yaml -collator: - mode: chat - include_thinking_in_loss: true - native_thinking_alignment: true - kd_alignment_regions: [reasoning, close, answer, eot] - num_packed_rows: 1 -``` - -Each row can supply `tools`. Set `use_preserving_dataset: true` for JSONL tool -data to retain heterogeneous argument objects without Arrow adding null fields. - -The final assistant turn supervises reasoning text (when -`include_thinking_in_loss` is true), the closing ``, answers, tool payloads, -and EOT. Opening thinking scaffolds, formatting-only tokens, schemas, -user/system/tool-result text, earlier turns, and padding -are excluded. `kd_alignment_regions` further restricts cross-tokenizer KD without -changing CE or same-tokenizer KD masks. Prose and each tool call align separately; -permitted Boolean/null spelling differences retain CE but omit the unequal piece -from exact-text KD. Missing or transformed required content fails explicitly. -An empty or whitespace-only final assistant turn is rejected; tool-only final -turns are supported. Supervising a separate `reasoning_content` field requires native -alignment. Ordinary chat supports reasoning embedded in `content` only when the -chat template renders that content verbatim. For Qwen-style templates that extract -inline `...` into a separate reasoning block, including the stock -Qwen3 template and `qwen3_history.jinja`, set -`collator.native_thinking_alignment=true`. Ordinary chat rejects requests to -supervise a separate reasoning field without native alignment. - -Qwen3-4B's stock template drops historical reasoning when a later user turn -follows it. To retain that reasoning as context, apply the maintained override to -**each Qwen tokenizer** used by the collator (both `policy.tokenizer` and the -appropriate `teachers[i].tokenizer`): - -```yaml -tokenizer: - name: Qwen/Qwen3-4B - tokenizer_kwargs: - revision: 1cfa9a7208912126459214e8b04321603b3df60c - chat_template: examples/chat_templates/qwen3_history.jinja - chat_template_kwargs: - enable_thinking: true - preserve_thinking: true - truncate_history_thinking: false -``` - -The template accepts explicit `reasoning_content` and leading inline -`...` content. An explicit `truncate_history_thinking` wins: -`false` retains history and `true` uses stock history truncation. When absent, -`preserve_thinking=true` enables retention; otherwise stock behavior applies. -`enable_thinking` preserves Qwen3's generation-prompt behavior, including its -empty thinking block when false; it does not delete supplied reasoning. Training -renders with `add_generation_prompt=False`. Native training requires the final -turn's supervised reasoning to be retained and rejects a template that silently -drops it. Historical reasoning follows the template's context-retention policy -and does not contribute loss. - -The override adds explicit history-retention controls to the pinned Qwen3-4B -template while preserving its tool serialization and generation behavior. The shared offset matcher compares literal decoded strings; whitespace -differences are unequal. Native Unicode repair is separate and requires both -NFC-normalized decodings to equal the original source region. Ordinary chat uses -one canonical nonempty content key on both sides for the final assistant turn. -Content tokens that also absorb surrounding whitespace retain CE supervision; -their full decoded strings still determine exact-match correctness. - -The [`xtoken_off_policy_distillation.sh`](../../tests/functional/xtoken_off_policy_distillation.sh) -functional test covers raw text. The standalone -[`xtoken_native_chat.sh`](../../tests/functional/xtoken_native_chat.sh) functional -test covers native chat and is registered separately in the fast and full CI -suites. Its -[`chat configuration`](../../examples/configs/xtoken_native_chat.yaml) -uses a SmolLM2-135M-Instruct student and Qwen3-1.7B teacher with distinct -tokenizer vocabularies and token boundaries. Both use the history-preserving -ChatML override so the test retains reasoning, tool calls, and end-of-turn -supervision. The script reads the checked-in multi-turn arithmetic -[`JSONL fixture`](../../tests/functional/test_data/xtoken_native_chat.jsonl), builds -its projection matrix, and checks that three training steps produce finite -losses. It owns its experiment directory and logs, so it can run independently: - -```bash -bash tests/functional/xtoken_native_chat.sh -``` - -The chat scenario is functional coverage, not a nightly convergence benchmark. -DP alignment, native student-only SFT, sequence packing, and backend/loss changes -are outside this port. - -### Loss-mode knobs - -`loss_fn` has two flags that pick between three behaviors: - -| `gold_loss` | `xtoken_loss` | Behavior | -|---|---|---| -| `false` | (inert) | **P-KL** — full-vocab teacher logits; the loss derives a microbatch-global top-k inside, projects the student into teacher vocab via the projection matrix, and chunk-averages KL on the top-k subset. CE term is added. | -| `true` | `false` | **Gold loss** — split the vocab into an *exact-token-mapped* common set (KL) and an *uncommon* tail (sorted L1). | -| `true` | `true` | **H-KL (gold + xtoken)** — same as gold, but relax the exact-map threshold to `>= 0.6` and allow multi-token projections to count as exact maps via a collision-replacement rule. | - -Other relevant fields: - -- `loss_fn.temperature` — softmax temperature applied symmetrically to student and teacher logits before KL. -- `loss_fn.vocab_topk` — microbatch-global top-k size for the P-KL path and for same-tokenizer teachers, whatever `gold_loss` is (the all-same-tokenizer `averaged_logits` path uses the full vocab). -- `loss_fn.uncommon_topk` — cap on the L1 uncommon-tail sort in the gold path (defaults to 8192). -- `loss_fn.reverse_kl` — compute `KL(student || teacher)` instead of `KL(teacher || student)`. - -## Results — 100-step multi-teacher run - -```bash -uv run python examples/run_xtoken_off_policy_distillation.py \ - --config examples/configs/xtoken_multiteacher_off_policy_distillation.yaml \ - teachers.0.aligner.projection_matrix_path=cross_tokenizer_data/projection_matrix_llama_phi-mini_top4.pt \ - cluster.gpus_per_node=8 \ - cluster.num_nodes=1 -``` - -This run distills a `meta-llama/Llama-3.2-1B` student from two teachers at -once: `microsoft/Phi-4-mini-instruct` (a cross-tokenizer teacher, projected -through its projection matrix) and `meta-llama/Llama-3.2-3B` (which shares the -student's tokenizer, so it contributes a direct top-K KL (`loss_fn.vocab_topk`) -with no projection). The per-teacher objectives are summed (`loss_fn.kd_loss_mode=sum`). -Config: global batch 96, micro-batch 1, sequence -length 2048, 100 steps, 2 nodes (8 GPUs each), on the default -Nemotron-Pretraining-Specialized-v1.1 / Formal-Logic corpus. The distillation -objective converges and the student tracks both teachers more closely over -training: - -train/loss, train/kl_loss, train/ce_loss, and train/accuracy over 100 multi-teacher distillation steps - -- **Loss** falls from ≈1.51 to ≈0.70. -- **KL loss** (summed over both teachers) falls from ≈4.89 to ≈1.97. Almost - all of it comes from the cross-tokenizer Phi-4-mini teacher (≈4.75 → ≈1.76); - the same-tokenizer Llama-3.2-3B teacher's direct KL is already small and - stays there (≈0.14 → ≈0.21). -- **CE loss** falls from ≈0.75 to ≈0.35. -- **Top-1 accuracy** rises from ≈0.82 to ≈0.91. - -### Downstream evaluation - -Benchmark scores of the distilled student against the undistilled -`meta-llama/Llama-3.2-1B` base model: - -| Task | Base Llama-3.2-1B | Distilled | -|---|---|---| -| MMLU | 32.05 | 40.24 | -| GSM8K | 5.69 | 5.76 | - -### Throughput and memory - -Measured on the same run (per training step, micro-batch 1, sequence length -2048, with the two teachers forwarded serially): - -| Metric | Value | -|---|---| -| Mean step time | 6.65 s (min 6.22 s) | -| Teacher forward (both teachers) | 4.94 s mean — the dominant per-step cost | -| Training throughput | ≈29.5k valid tokens/s (196,512 tokens/step ÷ mean step time) | - -Each teacher's full-vocab logits stay on-node: the producer publishes a -rank-level `[B_r, T_t, V_t]` bf16 tray and hands the student a CUDA IPC handle -to it, so teacher logits never cross the network even with two teachers in -play. - -## Where files live - -| Stage | Tool | Default output | -|---|---|---| -| Build multi-token | `tools/x_token/minimal_projection_via_multitoken.py` | `/projection_map__to__multitoken_top__double[_special].pt` | -| Reapply exact map | `tools/x_token/reapply_exact_map.py` | `_exact_map_remapped.pt` | -| Sort and trim | `tools/x_token/sort_and_cut_projection_matrix.py` | `/_top__sorted[_preservelast].pt` (or `--output_path`) | -| Train | `examples/run_xtoken_off_policy_distillation.py` | per the run's `logger.log_dir` and `checkpointing.checkpoint_dir` | - -## Related - -Listed in the order they run during a training step: - -- Config exemplar: [`examples/configs/xtoken_off_policy_distillation.yaml`](../../examples/configs/xtoken_off_policy_distillation.yaml) -- Trainer module: `nemo_rl/algorithms/xtoken_off_policy_distillation.py` -- KD data processor: `nemo_rl/data/processors.py::kd_data_processor` -- Cross-tokenizer collator: `nemo_rl/data/cross_tokenizer/cross_tokenizer_collate.py::CrossTokenizerCollator` -- Token alignment: `nemo_rl/algorithms/x_token/token_aligner.py::TokenAligner` -- Loss implementation: `nemo_rl/algorithms/loss/loss_functions.py::CrossTokenizerDistillationLossFn` diff --git a/docs/index.md b/docs/index.md index 62cab9aedcd..9992b8745a4 100644 --- a/docs/index.md +++ b/docs/index.md @@ -171,10 +171,10 @@ Extend a model's context window with YaRN RoPE scaling on the Megatron backend f ::: :::{grid-item-card} {octicon}`git-compare` Cross-Tokenizer Distillation -:link: guides/xtoken-off-policy-distillation +:link: guides/cross-tokenizer-distillation :link-type: doc -Off-policy distillation across mismatched tokenizers — build a (student, teacher) projection matrix and run x-token KD via CUDA-IPC teacher logits. +Off-policy distillation across mismatched tokenizers with X-Token (projection matrix) or CSD (sub-token tables), via CUDA-IPC teacher logits. ::: :::{grid-item-card} {octicon}`arrow-both` Weight Refit @@ -335,7 +335,7 @@ guides/single-controller.md guides/quantization-aware-rl.md guides/eagle3-speculative-decoding.md guides/yarn-long-context.md -guides/xtoken-off-policy-distillation.md +guides/cross-tokenizer-distillation.md guides/refit.md guides/checkpoint-engine-refit.md guides/dynamo-generation.md @@ -389,6 +389,7 @@ design-docs/loss-functions.md design-docs/fsdp2-parallel-plan.md design-docs/training-backends.md design-docs/sequence-packing-and-dynamic-batching.md +design-docs/csd-lockstep-sequence-packing.md design-docs/env-vars.md design-docs/nemo-gym-integration.md design-docs/modelopt-real-quant-architecture.md diff --git a/examples/configs/csd.yaml b/examples/configs/csd.yaml new file mode 100644 index 00000000000..c5b3be5679b --- /dev/null +++ b/examples/configs/csd.yaml @@ -0,0 +1,354 @@ +# CSD cross-tokenizer distillation. Build both tokenizer tables first; see +# the CSD section of docs/guides/cross-tokenizer-distillation.md. +defaults: distillation_math.yaml +distillation: + num_prompts_per_step: 64 + max_num_steps: 5000 + max_num_epochs: 1 + seed: 42 + val_period: 0 + val_at_start: false + val_at_end: false + offload_student_after_step: true +loss_fn: + temperature: 1.0 + vocab_topk: 8192 + reverse_kl: false + kl_loss_weight: 1.0 # used only when dynamic_loss_scaling is false + ce_loss_scale: 0.1 # used only when dynamic_loss_scaling is false + dynamic_loss_scaling: true # rescale KD to CE's size each step; the weights above are ignored + common_indices_from_subtoks: true + kl_chunk_shift: true + csd_loss_fn: jsd + csd_last_pos_loss_fn: null + csd_jsd_beta: 0.5 + csd_mismatch_loss_beta: 5.0 + csd_noise_filter_topk: 0 + csd_realized_only_mismatch: false + teacher_topk_ipc_k: 128 # 0 sends dense teacher logits + teacher_topk_ipc_support_mode: row_topk + teacher_topk_ipc_keep_realized: true + csd_selected_row_chunk_size: 128 # owner-local Megatron path: student/teacher rows per tile + csd_support_chunk_batch_size: 16 # owner-local Megatron path: alignment chunks per batch + kd_loss_mode: sum + sum_weights_metric: null + alpha: 1.0 + normalize_teacher_by_vocab: false +checkpointing: + enabled: false + checkpoint_dir: checkpoints/csd + metric_name: train:loss + higher_is_better: false + keep_top_k: 3 + save_period: 600 + checkpoint_must_save_by: null + save_optimizer: true +policy: + model_name: meta-llama/Llama-3.2-1B + tokenizer: + name: meta-llama/Llama-3.2-1B + train_global_batch_size: 64 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 2048 + make_sequence_length_divisible_by: 1 + precision: bfloat16 + logprob_chunk_size: null + offload_optimizer_for_logprob: true + router_replay: + enabled: false + automodel_cfg: + enabled: false + checkpoint: + model_save_format: "safetensors" + # Canonical Automodel modes: "false", "final", or "every". + save_consolidated: "false" + single_rank_consolidation: false + consolidation_timeout_minutes: 30 + cpu_offload: false + sequence_parallel: false + activation_checkpointing: true + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + load_precision: float32 + shard_before_load: false + megatron_cfg: + enabled: true + empty_unused_memory_level: 0 + activation_checkpointing: false + recompute_granularity: full + recompute_modules: null + tensor_model_parallel_size: 1 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 1 + pipeline_dtype: ${policy.precision} + sequence_parallel: false + freeze_moe_router: true + moe_router_dtype: fp64 + moe_router_load_balancing_type: none + moe_router_bias_update_rate: 0.0 + moe_permute_fusion: true + apply_rope_fusion: true + bias_activation_fusion: true + defer_fp32_logits: false + moe_per_layer_logging: false + moe_enable_deepep: false + moe_token_dispatcher_type: alltoall + moe_shared_expert_overlap: false + use_gloo_process_groups: false + gradient_accumulation_fusion: false + use_fused_weighted_squared_relu: false + checkpoint: + async_save: false + ckpt_assume_constant_structure: false + force_reconvert_from_hf: false + optimizer: + optimizer: adam + lr: 5.0e-05 + min_lr: 0.0 + weight_decay: 0.001 + bf16: true + fp16: false + params_dtype: float32 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1e-8 + sgd_momentum: 0.9 + use_distributed_optimizer: true + use_precision_aware_optimizer: true + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + clip_grad: ${policy.max_grad_norm} + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: constant + lr_decay_style: cosine + lr_decay_iters: 5000 + lr_warmup_iters: 250 + lr_warmup_init: 1.0e-06 + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: true + overlap_param_gather: true + use_custom_fsdp: false + data_parallel_sharding_strategy: optim_grads_params + fp8_cfg: + enabled: false + fp8: e4m3 + fp8_recipe: blockwise + fp8_param: false + max_grad_norm: 1.0 + freeze_parameter_patterns: null + dynamic_batching: + enabled: false + train_mb_tokens: 2048 + logprob_mb_tokens: 2048 + sequence_length_round: 64 + sequence_packing: + enabled: false + train_mb_tokens: 2048 + logprob_mb_tokens: 2048 + algorithm: modified_first_fit_decreasing + microbatch_order: packer + sequence_length_round: 64 + fuse_loss: false + optimizer: + name: torch.optim.AdamW + kwargs: + lr: 5.0e-05 + weight_decay: 0.001 + betas: + - 0.9 + - 0.95 + eps: 1e-8 + foreach: false + fused: false + scheduler: + - name: torch.optim.lr_scheduler.LinearLR + kwargs: + start_factor: 0.02 + end_factor: 1.0 + total_iters: 250 + - name: torch.optim.lr_scheduler.CosineAnnealingLR + kwargs: + T_max: 4750 + eta_min: 0.0 + - milestones: + - 250 +teacher: null +teachers: +- model_name: Qwen/Qwen3-4B + aligner: + projection_matrix_path: null + drop_first_assistant_chunk_kl: false + student_to_teacher_table_path: cross_tokenizer_data/csd/student_to_teacher.pt + teacher_to_student_table_path: cross_tokenizer_data/csd/teacher_to_student.pt + weight: 1.0 + tokenizer: + name: Qwen/Qwen3-4B + train_global_batch_size: 64 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 2048 + make_sequence_length_divisible_by: 1 + precision: bfloat16 + logprob_chunk_size: null + offload_optimizer_for_logprob: false + offload_policy_before_refit: false + offload_optimizer_for_refit: true + automodel_cfg: + enabled: false + checkpoint: + model_save_format: "safetensors" + save_consolidated: "false" + cpu_offload: false + sequence_parallel: false + activation_checkpointing: true + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + load_precision: float32 + shard_before_load: false + megatron_cfg: + enabled: true + empty_unused_memory_level: 0 + activation_checkpointing: false + recompute_granularity: full + recompute_modules: null + tensor_model_parallel_size: 1 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 1 + pipeline_dtype: bfloat16 + sequence_parallel: false + freeze_moe_router: true + moe_router_dtype: fp64 + moe_router_load_balancing_type: none + moe_router_bias_update_rate: 0.0 + moe_permute_fusion: true + apply_rope_fusion: true + bias_activation_fusion: true + defer_fp32_logits: false + moe_per_layer_logging: false + moe_enable_deepep: false + moe_token_dispatcher_type: alltoall + moe_shared_expert_overlap: false + use_gloo_process_groups: false + gradient_accumulation_fusion: false + use_fused_weighted_squared_relu: false + checkpoint: + async_save: false + ckpt_assume_constant_structure: false + force_reconvert_from_hf: false + optimizer: + optimizer: adam + lr: 5.0e-05 + min_lr: 5.0e-05 + weight_decay: 0.001 + bf16: true + fp16: false + params_dtype: float32 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1e-8 + sgd_momentum: 0.9 + use_distributed_optimizer: true + use_precision_aware_optimizer: true + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + clip_grad: 1.0 + scheduler: + start_weight_decay: 0.001 + end_weight_decay: 0.001 + weight_decay_incr_style: constant + lr_decay_style: constant + lr_decay_iters: 1000 + lr_warmup_iters: 10 + lr_warmup_init: 5.0e-06 + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: true + overlap_param_gather: true + use_custom_fsdp: false + data_parallel_sharding_strategy: optim_grads_params + fp8_cfg: + enabled: false + fp8: e4m3 + fp8_recipe: blockwise + fp8_param: false + max_grad_norm: 1.0 + dynamic_batching: + enabled: false + train_mb_tokens: 2048 + logprob_mb_tokens: 2048 + sequence_length_round: 64 + sequence_packing: + enabled: false + train_mb_tokens: 2048 + logprob_mb_tokens: 2048 + algorithm: modified_first_fit_decreasing + microbatch_order: packer + sequence_length_round: 64 + fuse_loss: false + optimizer: + name: torch.optim.AdamW + kwargs: + lr: 5.0e-05 + weight_decay: 0.001 + betas: + - 0.9 + - 0.95 + eps: 1e-8 + foreach: false + fused: false +data: + max_input_seq_length: 2048 + shuffle: true + num_workers: 4 + train: + dataset_name: arrow_text + processor: kd_data_processor + data_files: nvidia/Nemotron-Pretraining-Specialized-v1.1 + subset: Nemotron-Pretraining-Formal-Logic + split: train + text_key: text + characters_per_sample: null + validation: null +logger: + log_dir: logs/csd + num_val_samples_to_print: 0 + wandb_enabled: false + swanlab_enabled: false + mlflow_enabled: false + tensorboard_enabled: false + monitor_gpus: true + wandb: + project: csd + name: csd + swanlab: + project: csd + name: csd + tensorboard: + log_dir: tb_logs-csd + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 +cluster: + gpus_per_node: 8 + num_nodes: 1 +algorithm: csd +collator: + mode: text + include_thinking_in_loss: false + native_thinking_alignment: false + kd_alignment_regions: null + num_packed_rows: 1 diff --git a/examples/configs/csd_multiteacher.yaml b/examples/configs/csd_multiteacher.yaml new file mode 100644 index 00000000000..4314661595c --- /dev/null +++ b/examples/configs/csd_multiteacher.yaml @@ -0,0 +1,269 @@ +# CSD cross-tokenizer distillation. Build both tokenizer tables first; see +# the CSD section of docs/guides/cross-tokenizer-distillation.md. +defaults: distillation_math.yaml +distillation: + num_prompts_per_step: 64 + max_num_steps: 5000 + max_num_epochs: 1 + seed: 42 + val_period: 0 + val_at_start: false + val_at_end: false + offload_student_after_step: false +loss_fn: + temperature: 1.0 + vocab_topk: 8192 + reverse_kl: false + kl_loss_weight: 1.0 # used only when dynamic_loss_scaling is false + ce_loss_scale: 0.1 # used only when dynamic_loss_scaling is false + dynamic_loss_scaling: true # rescale KD to CE's size each step; the weights above are ignored + common_indices_from_subtoks: true + kl_chunk_shift: true + csd_loss_fn: jsd + csd_last_pos_loss_fn: null + csd_jsd_beta: 0.5 + csd_mismatch_loss_beta: 5.0 + csd_noise_filter_topk: 0 + csd_realized_only_mismatch: false + teacher_topk_ipc_k: 128 # 0 sends dense teacher logits + teacher_topk_ipc_support_mode: row_topk + teacher_topk_ipc_keep_realized: true + csd_selected_row_chunk_size: 128 # owner-local Megatron path: student/teacher rows per tile + csd_support_chunk_batch_size: 16 # owner-local Megatron path: alignment chunks per batch + kd_loss_mode: sum + sum_weights_metric: null + alpha: 1.0 + normalize_teacher_by_vocab: false +checkpointing: + enabled: false + checkpoint_dir: checkpoints/csd + metric_name: train:loss + higher_is_better: false + keep_top_k: 3 + save_period: 600 + checkpoint_must_save_by: null + save_optimizer: true +policy: + model_name: meta-llama/Llama-3.2-1B + tokenizer: + name: meta-llama/Llama-3.2-1B + train_global_batch_size: 64 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 2048 + make_sequence_length_divisible_by: 1 + precision: bfloat16 + logprob_chunk_size: null + offload_optimizer_for_logprob: true + automodel_cfg: + enabled: true + checkpoint: + model_save_format: "safetensors" + # Canonical Automodel modes: "false", "final", or "every". + save_consolidated: "false" + single_rank_consolidation: false + consolidation_timeout_minutes: 30 + cpu_offload: false + sequence_parallel: false + activation_checkpointing: true + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + load_precision: float32 + shard_before_load: false + max_grad_norm: 1.0 + freeze_parameter_patterns: null + dynamic_batching: + enabled: false + train_mb_tokens: 2048 + logprob_mb_tokens: 2048 + sequence_length_round: 64 + sequence_packing: + enabled: false + train_mb_tokens: 2048 + logprob_mb_tokens: 2048 + algorithm: modified_first_fit_decreasing + microbatch_order: packer + sequence_length_round: 64 + fuse_loss: false + optimizer: + name: torch.optim.AdamW + kwargs: + lr: 5.0e-05 + weight_decay: 0.001 + betas: + - 0.9 + - 0.95 + eps: 1e-8 + foreach: false + fused: false + scheduler: + - name: torch.optim.lr_scheduler.LinearLR + kwargs: + start_factor: 0.02 + end_factor: 1.0 + total_iters: 250 + - name: torch.optim.lr_scheduler.CosineAnnealingLR + kwargs: + T_max: 4750 + eta_min: 0.0 + - milestones: + - 250 +teacher: null +teachers: +- model_name: microsoft/Phi-4-mini-instruct + aligner: + projection_matrix_path: null + drop_first_assistant_chunk_kl: false + student_to_teacher_table_path: cross_tokenizer_data/csd/phi4mini/student_to_teacher.pt + teacher_to_student_table_path: cross_tokenizer_data/csd/phi4mini/teacher_to_student.pt + weight: 1.0 + tokenizer: + name: microsoft/Phi-4-mini-instruct + train_global_batch_size: 64 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 2048 + make_sequence_length_divisible_by: 1 + precision: bfloat16 + logprob_chunk_size: null + offload_optimizer_for_logprob: false + offload_policy_before_refit: false + offload_optimizer_for_refit: true + automodel_cfg: + enabled: true + checkpoint: + model_save_format: "safetensors" + save_consolidated: "false" + cpu_offload: false + sequence_parallel: false + activation_checkpointing: true + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + load_precision: bfloat16 + shard_before_load: false + max_grad_norm: 1.0 + dynamic_batching: + enabled: false + train_mb_tokens: 2048 + logprob_mb_tokens: 2048 + sequence_length_round: 64 + sequence_packing: + enabled: false + train_mb_tokens: 2048 + logprob_mb_tokens: 2048 + algorithm: modified_first_fit_decreasing + microbatch_order: packer + sequence_length_round: 64 + fuse_loss: false + optimizer: + name: torch.optim.AdamW + kwargs: + lr: 5.0e-05 + weight_decay: 0.001 + betas: + - 0.9 + - 0.95 + eps: 1e-8 + foreach: false + fused: false +- model_name: meta-llama/Llama-3.2-3B + aligner: + projection_matrix_path: null + drop_first_assistant_chunk_kl: false + student_to_teacher_table_path: null + teacher_to_student_table_path: null + weight: 1.0 + tokenizer: + name: meta-llama/Llama-3.2-3B + train_global_batch_size: 64 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 2048 + make_sequence_length_divisible_by: 1 + precision: bfloat16 + logprob_chunk_size: null + offload_optimizer_for_logprob: false + offload_policy_before_refit: false + offload_optimizer_for_refit: true + automodel_cfg: + enabled: true + checkpoint: + model_save_format: "safetensors" + save_consolidated: "false" + cpu_offload: false + sequence_parallel: false + activation_checkpointing: true + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + load_precision: bfloat16 + shard_before_load: false + max_grad_norm: 1.0 + dynamic_batching: + enabled: false + train_mb_tokens: 2048 + logprob_mb_tokens: 2048 + sequence_length_round: 64 + sequence_packing: + enabled: false + train_mb_tokens: 2048 + logprob_mb_tokens: 2048 + algorithm: modified_first_fit_decreasing + microbatch_order: packer + sequence_length_round: 64 + fuse_loss: false + optimizer: + name: torch.optim.AdamW + kwargs: + lr: 5.0e-05 + weight_decay: 0.001 + betas: + - 0.9 + - 0.95 + eps: 1e-8 + foreach: false + fused: false +data: + max_input_seq_length: 2048 + shuffle: true + num_workers: 4 + train: + dataset_name: arrow_text + processor: kd_data_processor + data_files: nvidia/Nemotron-Pretraining-Specialized-v1.1 + subset: Nemotron-Pretraining-Formal-Logic + split: train + text_key: text + characters_per_sample: null + validation: null +logger: + log_dir: logs/csd + num_val_samples_to_print: 0 + wandb_enabled: false + swanlab_enabled: false + mlflow_enabled: false + tensorboard_enabled: false + monitor_gpus: true + wandb: + project: csd + name: csd + swanlab: + project: csd + name: csd + tensorboard: + log_dir: tb_logs-csd + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 +cluster: + gpus_per_node: 8 + num_nodes: 1 +algorithm: csd +collator: + mode: text + include_thinking_in_loss: false + native_thinking_alignment: false + kd_alignment_regions: null + num_packed_rows: 1 diff --git a/examples/configs/recipes/llm/csd-multiteacher-qwen3-4b-llama3.2-3b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.yaml b/examples/configs/recipes/llm/csd-multiteacher-qwen3-4b-llama3.2-3b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.yaml new file mode 100644 index 00000000000..5815098f2c9 --- /dev/null +++ b/examples/configs/recipes/llm/csd-multiteacher-qwen3-4b-llama3.2-3b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.yaml @@ -0,0 +1,142 @@ +defaults: ../../csd.yaml +distillation: + num_prompts_per_step: 8 + max_num_steps: 20 + val_period: 10 + val_at_start: true +data: + train: + split_validation_size: 64 + seed: 42 +policy: + train_global_batch_size: 8 + max_total_sequence_length: 1024 + make_sequence_length_divisible_by: 2 + automodel_cfg: + enabled: true + tensor_parallel_size: 4 + context_parallel_size: 2 + megatron_cfg: + enabled: false + scheduler: + - name: torch.optim.lr_scheduler.LinearLR + kwargs: + start_factor: 0.1 + end_factor: 1.0 + total_iters: 5 + - name: torch.optim.lr_scheduler.ConstantLR + kwargs: + factor: 1.0 + total_iters: 10000000000 + - milestones: + - 5 +teachers: +- model_name: Qwen/Qwen3-4B + aligner: + projection_matrix_path: null + drop_first_assistant_chunk_kl: false + student_to_teacher_table_path: null + teacher_to_student_table_path: null + weight: 1.0 + tokenizer: + name: Qwen/Qwen3-4B + train_global_batch_size: 8 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 1024 + make_sequence_length_divisible_by: 2 + precision: bfloat16 + logprob_chunk_size: null + offload_optimizer_for_logprob: false + automodel_cfg: + enabled: true + checkpoint: + model_save_format: "safetensors" + save_consolidated: "false" + cpu_offload: false + sequence_parallel: false + activation_checkpointing: true + tensor_parallel_size: 2 + context_parallel_size: 2 + custom_parallel_plan: null + load_precision: bfloat16 + max_grad_norm: 1.0 + dynamic_batching: + enabled: false + train_mb_tokens: 1024 + logprob_mb_tokens: 1024 + sequence_length_round: 64 + sequence_packing: + enabled: false + train_mb_tokens: 1024 + logprob_mb_tokens: 1024 + algorithm: modified_first_fit_decreasing + sequence_length_round: 64 + optimizer: + name: torch.optim.AdamW + kwargs: + lr: 5.0e-05 + weight_decay: 0.001 + betas: + - 0.9 + - 0.95 + eps: 1e-8 + foreach: false + fused: false +- model_name: meta-llama/Llama-3.2-3B + aligner: + projection_matrix_path: null + drop_first_assistant_chunk_kl: false + student_to_teacher_table_path: null + teacher_to_student_table_path: null + weight: 1.0 + tokenizer: + name: meta-llama/Llama-3.2-3B + train_global_batch_size: 8 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 1024 + make_sequence_length_divisible_by: 2 + precision: bfloat16 + logprob_chunk_size: null + offload_optimizer_for_logprob: false + automodel_cfg: + enabled: true + checkpoint: + model_save_format: "safetensors" + save_consolidated: "false" + cpu_offload: false + sequence_parallel: false + activation_checkpointing: true + tensor_parallel_size: 2 + context_parallel_size: 2 + custom_parallel_plan: null + load_precision: bfloat16 + max_grad_norm: 1.0 + dynamic_batching: + enabled: false + train_mb_tokens: 1024 + logprob_mb_tokens: 1024 + sequence_length_round: 64 + sequence_packing: + enabled: false + train_mb_tokens: 1024 + logprob_mb_tokens: 1024 + algorithm: modified_first_fit_decreasing + sequence_length_round: 64 + optimizer: + name: torch.optim.AdamW + kwargs: + lr: 5.0e-05 + weight_decay: 0.001 + betas: + - 0.9 + - 0.95 + eps: 1e-8 + foreach: false + fused: false +logger: + log_dir: logs/csd-multiteacher-qwen3-4b-llama3.2-3b-to-llama3.2-1b-1n8g-dtensor-tp4cp2 + wandb: + project: nemo-rl + name: csd-multiteacher-qwen3-4b-llama3.2-3b-to-llama3.2-1b-1n8g-dtensor-tp4cp2 diff --git a/examples/configs/recipes/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.yaml b/examples/configs/recipes/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.yaml new file mode 100644 index 00000000000..29fb43c5597 --- /dev/null +++ b/examples/configs/recipes/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.yaml @@ -0,0 +1,89 @@ +defaults: ../../csd.yaml +distillation: + num_prompts_per_step: 8 + max_num_steps: 20 + val_period: 10 + val_at_start: true +data: + train: + split_validation_size: 64 + seed: 42 +policy: + train_global_batch_size: 8 + max_total_sequence_length: 1024 + make_sequence_length_divisible_by: 2 + automodel_cfg: + enabled: true + tensor_parallel_size: 4 + context_parallel_size: 2 + megatron_cfg: + enabled: false + scheduler: + - name: torch.optim.lr_scheduler.LinearLR + kwargs: + start_factor: 0.1 + end_factor: 1.0 + total_iters: 5 + - name: torch.optim.lr_scheduler.ConstantLR + kwargs: + factor: 1.0 + total_iters: 10000000000 + - milestones: + - 5 +teachers: +- model_name: Qwen/Qwen3-4B + aligner: + projection_matrix_path: null + drop_first_assistant_chunk_kl: false + student_to_teacher_table_path: null + teacher_to_student_table_path: null + weight: 1.0 + tokenizer: + name: Qwen/Qwen3-4B + train_global_batch_size: 8 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 1024 + make_sequence_length_divisible_by: 2 + precision: bfloat16 + logprob_chunk_size: null + offload_optimizer_for_logprob: false + automodel_cfg: + enabled: true + checkpoint: + model_save_format: "safetensors" + save_consolidated: "false" + cpu_offload: false + sequence_parallel: false + activation_checkpointing: true + tensor_parallel_size: 2 + context_parallel_size: 2 + custom_parallel_plan: null + max_grad_norm: 1.0 + dynamic_batching: + enabled: false + train_mb_tokens: 1024 + logprob_mb_tokens: 1024 + sequence_length_round: 64 + sequence_packing: + enabled: false + train_mb_tokens: 1024 + logprob_mb_tokens: 1024 + algorithm: modified_first_fit_decreasing + sequence_length_round: 64 + optimizer: + name: torch.optim.AdamW + kwargs: + lr: 5.0e-05 + weight_decay: 0.001 + betas: + - 0.9 + - 0.95 + eps: 1e-8 + foreach: false + fused: false +logger: + log_dir: logs/csd-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2 + wandb: + project: nemo-rl + name: csd-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2 diff --git a/examples/configs/recipes/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-megatron-tp2cp2.yaml b/examples/configs/recipes/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-megatron-tp2cp2.yaml new file mode 100644 index 00000000000..4de959c1ba3 --- /dev/null +++ b/examples/configs/recipes/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-megatron-tp2cp2.yaml @@ -0,0 +1,153 @@ +defaults: ../../csd.yaml +distillation: + num_prompts_per_step: 8 + max_num_steps: 20 + val_period: 10 + val_at_start: true +data: + train: + split_validation_size: 64 + seed: 42 +policy: + train_global_batch_size: 8 + max_total_sequence_length: 1024 + make_sequence_length_divisible_by: 8 + megatron_cfg: + tensor_model_parallel_size: 2 + context_parallel_size: 2 + scheduler: + lr_decay_style: constant + lr_warmup_iters: 5 +teachers: +- model_name: Qwen/Qwen3-4B + aligner: + projection_matrix_path: null + drop_first_assistant_chunk_kl: false + student_to_teacher_table_path: cross_tokenizer_data/csd/student_to_teacher.pt + teacher_to_student_table_path: cross_tokenizer_data/csd/teacher_to_student.pt + weight: 1.0 + tokenizer: + name: Qwen/Qwen3-4B + train_global_batch_size: 8 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 1024 + make_sequence_length_divisible_by: 8 + precision: bfloat16 + logprob_chunk_size: null + offload_optimizer_for_logprob: false + offload_policy_before_refit: false + offload_optimizer_for_refit: true + automodel_cfg: + enabled: false + checkpoint: + model_save_format: "safetensors" + save_consolidated: "false" + cpu_offload: false + sequence_parallel: false + activation_checkpointing: true + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + load_precision: float32 + shard_before_load: false + megatron_cfg: + enabled: true + empty_unused_memory_level: 0 + activation_checkpointing: false + recompute_granularity: full + recompute_modules: null + tensor_model_parallel_size: 2 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 2 + pipeline_dtype: bfloat16 + sequence_parallel: false + freeze_moe_router: true + moe_router_dtype: fp64 + moe_router_load_balancing_type: none + moe_router_bias_update_rate: 0.0 + moe_permute_fusion: true + apply_rope_fusion: true + bias_activation_fusion: true + defer_fp32_logits: false + moe_per_layer_logging: false + moe_enable_deepep: false + moe_token_dispatcher_type: alltoall + moe_shared_expert_overlap: false + use_gloo_process_groups: false + gradient_accumulation_fusion: false + use_fused_weighted_squared_relu: false + checkpoint: + async_save: false + ckpt_assume_constant_structure: false + force_reconvert_from_hf: false + optimizer: + optimizer: adam + lr: 5.0e-05 + min_lr: 5.0e-05 + weight_decay: 0.001 + bf16: true + fp16: false + params_dtype: float32 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1e-8 + sgd_momentum: 0.9 + use_distributed_optimizer: true + use_precision_aware_optimizer: true + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + clip_grad: 1.0 + scheduler: + start_weight_decay: 0.001 + end_weight_decay: 0.001 + weight_decay_incr_style: constant + lr_decay_style: constant + lr_decay_iters: 1000 + lr_warmup_iters: 10 + lr_warmup_init: 5.0e-06 + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: true + overlap_param_gather: true + use_custom_fsdp: false + data_parallel_sharding_strategy: optim_grads_params + fp8_cfg: + enabled: false + fp8: e4m3 + fp8_recipe: blockwise + fp8_param: false + max_grad_norm: 1.0 + dynamic_batching: + enabled: false + train_mb_tokens: 2048 + logprob_mb_tokens: 2048 + sequence_length_round: 64 + sequence_packing: + enabled: false + train_mb_tokens: 2048 + logprob_mb_tokens: 2048 + algorithm: modified_first_fit_decreasing + microbatch_order: packer + sequence_length_round: 64 + fuse_loss: false + optimizer: + name: torch.optim.AdamW + kwargs: + lr: 5.0e-05 + weight_decay: 0.001 + betas: + - 0.9 + - 0.95 + eps: 1e-8 + foreach: false + fused: false +logger: + log_dir: logs/csd-qwen3-4b-to-llama3.2-1b-1n8g-megatron-tp2cp2 + wandb: + project: nemo-rl + name: csd-qwen3-4b-to-llama3.2-1b-1n8g-megatron-tp2cp2 diff --git a/examples/configs/xtoken_multiteacher_off_policy_distillation.yaml b/examples/configs/xtoken_multiteacher_off_policy_distillation.yaml index d5a655eae7a..87b8e66332d 100644 --- a/examples/configs/xtoken_multiteacher_off_policy_distillation.yaml +++ b/examples/configs/xtoken_multiteacher_off_policy_distillation.yaml @@ -25,7 +25,7 @@ loss_fn: # NOTE: the per-(student, teacher) projection matrix path now lives at # `teachers[i].aligner.projection_matrix_path` (see below), not here. The # knobs in this block are shared across all teachers. - # See docs/guides/xtoken-off-policy-distillation.md for how to produce the + # See docs/guides/cross-tokenizer-distillation.md for how to produce the # projection file. # Loss-mode selection (shared default; override per teacher via # teachers[i].gold_loss / .xtoken_loss in `sum` mode): diff --git a/examples/configs/xtoken_off_policy_distillation.yaml b/examples/configs/xtoken_off_policy_distillation.yaml index 1d5cd63cd75..2309fecc077 100644 --- a/examples/configs/xtoken_off_policy_distillation.yaml +++ b/examples/configs/xtoken_off_policy_distillation.yaml @@ -25,7 +25,7 @@ loss_fn: # NOTE: the per-(student, teacher) projection matrix path now lives at # `teachers[i].aligner.projection_matrix_path` (see below), not here. The # knobs in this block are shared across all teachers. - # See docs/guides/xtoken-off-policy-distillation.md for how to produce the + # See docs/guides/cross-tokenizer-distillation.md for how to produce the # projection file. # Loss-mode selection (shared default; override per teacher via # teachers[i].gold_loss / .xtoken_loss in `sum` mode): diff --git a/examples/run_csd.py b/examples/run_csd.py new file mode 100644 index 00000000000..1aeed916ee3 --- /dev/null +++ b/examples/run_csd.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Launch CSD using the shared cross-tokenizer distillation controller.""" + +from run_xtoken_off_policy_distillation import main + +if __name__ == "__main__": + main(default_config="csd.yaml") diff --git a/examples/run_xtoken_off_policy_distillation.py b/examples/run_xtoken_off_policy_distillation.py index 35655f880aa..0d7298b884b 100644 --- a/examples/run_xtoken_off_policy_distillation.py +++ b/examples/run_xtoken_off_policy_distillation.py @@ -18,13 +18,17 @@ import argparse import os import pprint +from pathlib import Path from omegaconf import OmegaConf from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.algorithms.xtoken_off_policy_distillation import ( + CSDMasterConfig, MasterConfig, setup, + validate_csd_packing_setup, + validate_csd_tokenizer_reuse, xtoken_off_policy_distillation_train, ) from nemo_rl.data.utils import setup_response_data @@ -46,19 +50,23 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]: parser.add_argument( "--config", type=str, default=None, help="Path to YAML config file" ) + parser.add_argument( + "--resolved-config-output", + type=str, + default=None, + help="Optional path for the fully resolved, validated configuration", + ) args, overrides = parser.parse_known_args() return args, overrides -def main() -> None: +def main(default_config: str = "xtoken_off_policy_distillation.yaml") -> None: """Main entry point.""" register_omegaconf_resolvers() args, overrides = parse_args() if not args.config: - args.config = os.path.join( - os.path.dirname(__file__), "configs", "xtoken_off_policy_distillation.yaml" - ) + args.config = os.path.join(os.path.dirname(__file__), "configs", default_config) config = load_config(args.config) if overrides: @@ -67,13 +75,23 @@ def main() -> None: config = OmegaConf.to_container(config, resolve=True) check_outdated_config(config) - config = MasterConfig(**config) - # Per-teacher same-vocab vs cross-tokenizer is determined solely by - # `teachers[i].aligner.projection_matrix_path` (null => same-vocab direct - # KL; set => cross-tokenizer). The consistency check (a same-vocab teacher - # must actually share the student's vocab) needs the real tokenizers and - # lives in `setup()` — comparing tokenizer *names* here is wrong, e.g. - # Llama-3.2-3B and Llama-3.2-1B have different names but the same vocab. + config = ( + CSDMasterConfig(**config) + if config.get("algorithm", "xtoken") == "csd" + else MasterConfig(**config) + ) + validate_csd_packing_setup(config) + if args.resolved_config_output: + resolved_config_path = Path(args.resolved_config_output).resolve() + resolved_config_path.parent.mkdir(parents=True, exist_ok=True) + OmegaConf.save( + config=OmegaConf.create(config.model_dump(mode="json")), + f=resolved_config_path, + ) + print(f"Resolved config written to: {resolved_config_path}", flush=True) + # X-Token selects cross-tokenizer teachers through projection_matrix_path; + # CSD uses its sub-token table pair. CSD validates full tokenizer/template + # equivalence before allowing a same-tokenizer teacher to reuse student IDs. print("Final config:") pprint.pprint(config) @@ -85,8 +103,6 @@ def main() -> None: flush=True, ) - init_ray() - # Student tokenizer + one tokenizer per teacher (same-tokenizer teachers # included — needed for vocab sizing in setup()). student_tokenizer = get_tokenizer(config.policy["tokenizer"]) @@ -94,6 +110,8 @@ def main() -> None: get_tokenizer(teacher.policy_config()["tokenizer"]) for teacher in config.teachers ] + if config.algorithm == "csd": + validate_csd_tokenizer_reuse(config, student_tokenizer, teacher_tokenizers) # `env_configs=None` skips the env-creation block (no rollout path); # `setup_response_data` then handles dataset construction, the optional @@ -104,6 +122,8 @@ def main() -> None: student_tokenizer, config.data, env_configs=None ) + init_ray() + ( student_policy, teacher_policies, diff --git a/nemo_rl/algorithms/loss/csd.py b/nemo_rl/algorithms/loss/csd.py new file mode 100644 index 00000000000..592abaeeff3 --- /dev/null +++ b/nemo_rl/algorithms/loss/csd.py @@ -0,0 +1,4591 @@ +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import collections +import os +from typing import Annotated, Any, Literal, Mapping, NotRequired, Optional, TypedDict + +import numpy as np +import torch +from pydantic import BaseModel, Field, model_validator +from torch.distributed.tensor import DTensor + +from nemo_rl.algorithms.loss.interfaces import LossFunction, LossInputType, LossType +from nemo_rl.algorithms.utils import masked_mean +from nemo_rl.algorithms.x_token.loss_utils import ( + DenseTeacherRowReader, + LocalizedAlignment, + SparseTeacherLogits, + SparseTeacherRowReader, + TeacherRowPayload, + build_exact_token_map, + ce_label_mask, + cp_load_balanced_global_positions, + next_token_accuracy, + same_tokenizer_full_vocab_kl, + same_tokenizer_topk_kl, + student_next_token_ce, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.model_utils import ( + allgather_cp_contiguous_tensor, + chunked_vocab_parallel_logsumexp, + distributed_selected_logprobs, + get_next_token_logprobs_from_logits, + group_all_reduce_sum, + group_all_reduce_sum_with_grad, + group_all_reduce_sum_with_grad_backward_sum, + to_local_if_dtensor, + vocab_parallel_argmax, +) + +CSD_STUDENT_LOSS_CHUNK_SIZE = 256 + + +def _generalized_jsd( + student_log_probs: torch.Tensor, + teacher_log_probs: torch.Tensor, + beta: float, +) -> torch.Tensor: + """Per-element generalized Jensen-Shannon Divergence on log-probabilities. + + Mirrors TRL's ``GOLDTrainer.generalized_jsd_loss`` (huggingface/trl @v1.0.0 + ``trl/experimental/gold/gold_trainer.py``). Inputs are log-probs (may be + sub-distribution slices - no normalization is applied here, so the + divergence is on raw probability mass). + + beta semantics (matches TRL): + * beta = 0 -> forward KL : KL(teacher || student) (mean-seeking) + * beta = 1 -> reverse KL : KL(student || teacher) (mode-seeking; TRL's default) + * beta = 0.5 -> symmetric JSD + * 0 < beta < 1 -> blended KL on a (1-beta).student + beta.teacher mixture + + Numerics: computation is forced to fp32 even when inputs are bf16/fp16. + bf16's ~7.8e-3 epsilon causes (student_log_probs - teacher_log_probs) to + underflow as the student converges, collapsing the matched-term loss to + exactly 0 mid-training. Cast back to the input dtype on return. + + Returns a tensor with the same shape as the inputs (per-element divergence, + matching F.kl_div(reduction='none')). Caller does the chunk/vocab sum. + """ + in_dtype = student_log_probs.dtype + s_lp = student_log_probs.to(torch.float32) + t_lp = teacher_log_probs.to(torch.float32) + + if beta == 0.0: + out = torch.nn.functional.kl_div( + s_lp, + t_lp, + reduction="none", + log_target=True, + ) + return out.to(in_dtype) + if beta == 1.0: + out = torch.nn.functional.kl_div( + t_lp, + s_lp, + reduction="none", + log_target=True, + ) + return out.to(in_dtype) + beta_t = torch.as_tensor(beta, dtype=torch.float32, device=s_lp.device) + # mixture_log_probs[v] = log( (1-beta)*p_s(v) + beta*p_t(v) ) + mixture_log_probs = torch.logsumexp( + torch.stack( + [ + s_lp + torch.log1p(-beta_t), + t_lp + torch.log(beta_t), + ] + ), + dim=0, + ) + kl_teacher = torch.nn.functional.kl_div( + mixture_log_probs, + t_lp, + reduction="none", + log_target=True, + ) + kl_student = torch.nn.functional.kl_div( + mixture_log_probs, + s_lp, + reduction="none", + log_target=True, + ) + out = beta_t * kl_teacher + (1 - beta_t) * kl_student + return out.to(in_dtype) + + +class CSDLossConfig(BaseModel, extra="allow"): + """CSD loss options and runtime-injected per-teacher metadata. + + Forward tables map student tokens to teacher sub-token chains; reverse + tables map teacher tokens to student chains. Optional projection matrices + provide the common-token mapping when ``common_indices_from_subtoks=False``. + """ + + temperature: Annotated[float, Field(gt=0)] = 1.0 + vocab_topk: Annotated[int, Field(gt=0)] = 8192 # Same-tokenizer direct KL only. + reverse_kl: bool = False + kl_loss_weight: float = 1.0 + ce_loss_scale: float = 0.1 + dynamic_loss_scaling: bool = True + kd_loss_mode: Literal["sum", "averaged_logits", "select_teacher"] = "sum" + normalize_teacher_by_vocab: bool = False + alpha: Annotated[float, Field(gt=0)] = 1.0 + sum_weights_metric: Optional[Literal["ce", "entropy", "max_prob"]] = None + + # Runtime-injected by setup; these are not user loss_fn fields. + student_vocab_size: Optional[int] = None + teacher_vocab_sizes: list[int] = Field(default_factory=list) + projection_matrix_paths: Optional[list[Optional[str]]] = None + teacher_weights: list[float] = Field(default_factory=list) + student_to_teacher_table_paths: Optional[list[Optional[str]]] = None + teacher_to_student_table_paths: Optional[list[Optional[str]]] = None + + # Shift each noninitial chunk to predictor positions before computing KD. + kl_chunk_shift: bool = True + # Build exact common-token pairs from forward tables instead of projections. + common_indices_from_subtoks: bool = False + csd_loss_fn: Literal["kl", "jsd", "bce"] = "jsd" + csd_jsd_beta: Annotated[float, Field(ge=0, le=1)] = 0.5 + # None inherits csd_loss_fn for the last-position partition. + csd_last_pos_loss_fn: Optional[Literal["kl", "jsd", "bce"]] = None + # Multiplier on each mismatch (non-1-to-1) chunk loss relative to a + # 1-to-1 chunk loss. + csd_mismatch_loss_beta: Annotated[float, Field(ge=0)] = 5.0 + # Zero disables filtering; otherwise retain realized labels in teacher top-k. + csd_noise_filter_topk: Annotated[int, Field(ge=0)] = 0 + # Compare only the realized chain probability and its complementary bucket. + csd_pure_alm: bool = False + # Score 1-to-many and many-to-1 chunks by their realized token chains + # against REST instead of prefix-support alternatives. + csd_realized_only_mismatch: bool = False + csd_alm_bce_tau: Annotated[float, Field(gt=0)] = 1.0 + # Positive k sends each teacher row's top-k logits plus its full-vocab logZ + # (the KD then compares top-k plus a REST bucket); zero sends dense logits. + teacher_topk_ipc_k: Annotated[int, Field(ge=0)] = 128 + teacher_topk_ipc_support_mode: Literal["row_topk"] = "row_topk" + teacher_topk_ipc_keep_realized: bool = True + # Owner-local evaluation (eligible Megatron students): selected student or + # teacher rows per tile, and alignment chunks evaluated per batch. Both + # bound memory only; they do not change the loss. + csd_selected_row_chunk_size: Annotated[int, Field(gt=0)] = 128 + csd_support_chunk_batch_size: Annotated[int, Field(gt=0)] = 16 + + # Runtime-injected by setup: evaluate cross-tokenizer KD on the context- + # parallel rank that owns each chunk and return a CP-partial loss. + owner_local: bool = False + + @model_validator(mode="before") + @classmethod + def _reject_removed_position_0_keys(cls, value: Any) -> Any: + if isinstance(value, dict): + for key in ( + "csd_position_0_kl", + "csd_mismatch_pos0_alpha", + "csd_mismatch_pos0_weight", + "csd_mismatch_loss_scale", + ): + if key in value: + raise ValueError( + f"loss_fn.{key} was removed: there is no position-0 " + "term, and mismatch chunks are weighted by " + "csd_mismatch_loss_beta alone." + ) + return value + + +def owner_local_config_violations(cfg: CSDLossConfig) -> list[str]: + """CSD loss settings that owner-local evaluation does not support. + + An empty list means the loss can run owner-local. Owner-local KD keeps a + static-weight sum of per-teacher terms, evaluates common chunks over a + multi-token support (so BCE, which needs a binary partition, is allowed + only for realized-only mismatch chunks), and with sparse teacher IPC + requires every request's realized label to be kept. + """ + violations: list[str] = [] + if cfg.kd_loss_mode != "sum": + violations.append(f"kd_loss_mode={cfg.kd_loss_mode!r} (needs 'sum')") + if cfg.sum_weights_metric is not None: + violations.append( + f"sum_weights_metric={cfg.sum_weights_metric!r} (needs static weights)" + ) + if cfg.csd_pure_alm: + violations.append("csd_pure_alm=true") + if cfg.csd_loss_fn == "bce": + violations.append("csd_loss_fn='bce' (common chunks are not binary)") + if cfg.csd_last_pos_loss_fn == "bce" and not cfg.csd_realized_only_mismatch: + violations.append( + "csd_last_pos_loss_fn='bce' without csd_realized_only_mismatch=true" + ) + if ( + cfg.teacher_topk_ipc_k > 1 + and not cfg.teacher_topk_ipc_keep_realized + and not cfg.csd_realized_only_mismatch + ): + violations.append( + "teacher_topk_ipc_keep_realized=false with teacher_topk_ipc_k > 1" + ) + return violations + + +class CSDLossDataDict(TypedDict): + """Student-side keys are fixed; teacher-side keys are teacher-indexed. + + Only the student keys below are static. Each teacher ``i`` contributes a + dynamic set of keys produced by ``CrossTokenizerCollator`` / the trainer and + so cannot be enumerated here: + + - Every teacher contributes exactly one logits transport. Same-tokenizer + teachers and dense cross-tokenizer teachers use + ``teacher_{i}_full_logits_ipc`` (List[B] of CUDA IPC handle dicts with + TP/CP shard metadata). Sparse cross-tokenizer teachers use + ``teacher_{i}_sparse_logits_ipc`` with row-top-k logits/token ids, the + exact full-vocabulary log-normalizer, and optional realized-label + membership. ``prepare_loss_input`` P2P-reads and reassembles either form. + - Cross-tokenizer teacher only: ``teacher_{i}_input_ids`` / + ``teacher_{i}_token_mask`` ``[B, T_t]`` and ``alignment_{i}_*`` + (``pair_valid`` / ``pair_is_correct`` ``[B, max_pairs]``; + ``student_chunk_id`` ``[B, T_s]``; ``teacher_chunk_id`` ``[B, T_t]``; + partition masks; ``num_chunks``). + - Same-tokenizer teacher: no ``teacher_{i}_input_ids`` / ``alignment_{i}_*``; + it reuses the student tokenization (identity 1:1 aligned). + """ + + input_ids: torch.Tensor + input_lengths: torch.Tensor + token_mask: torch.Tensor + kd_token_mask: NotRequired[torch.Tensor] + sample_mask: torch.Tensor + + +class CSDLossFn(LossFunction): + """CSD cross-tokenizer distillation loss. + + A cross-tokenizer teacher's KD term is computed by ``_compute_csd_loss``: a common-vocab partition KL/JSD + on the exact 1-to-1 token map plus, for non-1-to-1 chunks, a batched + mismatch partition KL over prefix-support alternatives (MtoM-as-ALM + rest + bucket) weighted by ``csd_mismatch_loss_beta``. It is combined with a + next-token student CE term as ``kl_loss_weight * kd + ce_loss_scale * ce`` — + or, when ``dynamic_loss_scaling`` is set, with the KD term rescaled each + step to match the detached CE magnitude. + + Multi-teacher: ``setup`` injects per-teacher metadata (projection paths, + weights, vocab sizes, pseudo-target-table paths). The per-teacher KD terms + are aggregated by ``kd_loss_mode`` (``sum`` / ``averaged_logits`` / + ``select_teacher``) and combined with a single student CE term. A teacher + without CSD tables or a projection is a *same-tokenizer* teacher: alignment + is skipped and its KD term is a direct top-k per-position KL on the + shared vocab (top-k selected student-side over the reassembled full-vocab + teacher logits). The single-teacher path is just ``num_teachers == 1``. + + Inputs (via ``LossInputType.DISTILLATION_CROSS_TOKENIZER``): + logits: ``[B, T_s, V_s]`` raw student logits from the worker forward. + student_logits_contig: CP-relaid contiguous student logits shared by + every teacher's KD term. + teacher_full_logits_by_idx: ``dict[int, [B, T, V_t]]`` full-vocab teacher + logits for dense teachers, rebuilt from CUDA IPC handles by + ``prepare_loss_input`` (see + :func:`nemo_rl.algorithms.x_token.loss_utils.rebuild_teacher_full_logits_from_ipc`). + teacher_sparse_logits_by_idx: Per-sparse-teacher row-top-k logits, + token ids, exact full-vocabulary log-normalizers, and optional + realized-label top-k membership, reconstructed for the full + teacher sequence. + teacher_row_payloads_by_idx: Owner-local evaluation only: each + cross-tokenizer teacher's IPC records, read row by row by the rank + that owns each chunk (see ``owner_local``). + aligns_by_idx: ``dict[int, LocalizedAlignment]`` per teacher (cross-tok: + localized chunk alignment; same-tok: thin, student fields only). + + Inputs (via ``data: BatchedDataDict``): + See :class:`CSDLossDataDict`. + + Returns: + ``(loss, metrics)``. Aggregate metrics: ``loss``, ``kl_loss`` (the + aggregated KD term), ``ce_loss``, ``kl_loss_scale``, ``accuracy``, + ``num_valid_samples``. Per-teacher metrics are suffixed ``_t{i}`` (e.g. + ``kl_loss_t0``, ``top1_acc_per_chunk_t0``, ``weight_t0``); + routing observability is reported as ``teacher_{i}/routed_samples``, + ``teacher_{i}/routed_tokens``, and ``teacher_{i}/weighted_kl``; + ``select_teacher`` additionally reports ``selected_teacher``. + """ + + # Group the global normalizers (_dp_global_masked_mean, the CSD + # global_valid_chunks) reduce over; set per call by __call__. Declared at + # class scope so it exists even for instances built without __init__ (the + # unit tests construct bare instances to exercise single methods). None => + # the legacy WORLD reduce; see loss_replica_group for why WORLD deadlocks + # under pipeline parallelism. + _dp_cp_group: Optional[torch.distributed.ProcessGroup] = None + # Whether per-window same-tokenizer KD terms are summed across CP; set per + # call by __call__ (see ``cp_normalized_by_caller``) like ``_dp_cp_group``. + _sum_window_kd_over_cp: bool = False + + loss_type = LossType.TOKEN_LEVEL + input_type = LossInputType.DISTILLATION_CROSS_TOKENIZER + + def __init__(self, cfg: CSDLossConfig | Mapping[str, Any]): + cfg = ( + cfg if isinstance(cfg, CSDLossConfig) else CSDLossConfig.model_validate(cfg) + ) + # Dynamic teacher weighting (sum_weights_metric) and normalize_teacher_by_vocab + # are only applied in kd_loss_mode="sum"; reject the combo instead of + # silently ignoring them under the other modes. + if cfg.sum_weights_metric is not None and cfg.kd_loss_mode != "sum": + raise ValueError( + f"sum_weights_metric={cfg.sum_weights_metric!r} is only applied " + f"in kd_loss_mode='sum'; it is ignored by '{cfg.kd_loss_mode}'. " + "Unset one of them." + ) + if cfg.normalize_teacher_by_vocab and cfg.kd_loss_mode != "sum": + raise ValueError( + "normalize_teacher_by_vocab is only applied in kd_loss_mode='sum'; " + f"it is ignored by '{cfg.kd_loss_mode}'. Unset one of them." + ) + # averaged_logits forms a convex combination of teacher logits + # (weight_i / sum(weights)); a zero weight-sum (all zeros, or a signed + # set cancelling to 0) makes that division undefined. Reject it here + # rather than fail with a deep ZeroDivisionError mid-step. + _weights = cfg.teacher_weights + if ( + cfg.kd_loss_mode == "averaged_logits" + and _weights is not None + and sum(_weights) == 0 + ): + raise ValueError( + "teacher_weights must not sum to zero in " + "kd_loss_mode='averaged_logits' (they form a convex combination " + f"of teacher logits); got teacher_weights={list(_weights)}." + ) + if cfg.student_vocab_size is None: + raise ValueError("CSD requires the runtime student_vocab_size") + if cfg.owner_local: + violations = owner_local_config_violations(cfg) + if violations: + raise ValueError( + "Owner-local CSD evaluation does not support " + + "; ".join(violations) + + "." + ) + # Evaluate cross-tokenizer KD on CP owner ranks (CP-partial loss). + self.owner_local = cfg.owner_local + # Global loss knobs (shared across all teachers). + self.temperature = cfg.temperature + # vocab_topk / reverse_kl are consumed by the same-vocab direct-KL path + # (same_tokenizer_topk_kl / same_tokenizer_full_vocab_kl); the + # cross-tokenizer path is CSD. + self.vocab_topk = cfg.vocab_topk + self.reverse_kl = cfg.reverse_kl + self.kl_loss_weight = cfg.kl_loss_weight + self.ce_loss_scale = cfg.ce_loss_scale + self.dynamic_loss_scaling = cfg.dynamic_loss_scaling + self.student_vocab_size = cfg.student_vocab_size + # Multi-teacher aggregation knobs. + self.kd_loss_mode = cfg.kd_loss_mode + self.normalize_teacher_by_vocab = cfg.normalize_teacher_by_vocab + self.alpha = cfg.alpha + # None selects the static per-teacher weights. + self.sum_weights_metric = cfg.sum_weights_metric + # Per-teacher metadata: parallel lists (one entry per ``teachers[i]``), + # injected by ``xtoken_off_policy_distillation.setup``. Every teacher ships + # full-vocab logits (the loss derives the top-k subset student-side), so + # there is no per-teacher ``send_full_logits`` flag. + self.projection_matrix_paths: list[Optional[str]] = ( + list(cfg.projection_matrix_paths) + if cfg.projection_matrix_paths is not None + else [None for _ in cfg.teacher_vocab_sizes] + ) + self.teacher_vocab_sizes = list(cfg.teacher_vocab_sizes) + self.teacher_weights = list(cfg.teacher_weights) + # Every per-teacher list must have the same length (one entry per + # teacher); a mismatch would otherwise surface as a deep IndexError + # mid-training instead of a clear error here. + per_teacher_lens = { + "projection_matrix_paths": len(self.projection_matrix_paths), + "teacher_vocab_sizes": len(self.teacher_vocab_sizes), + "teacher_weights": len(self.teacher_weights), + } + if len(set(per_teacher_lens.values())) != 1: + raise ValueError( + f"per-teacher lists must be equal length, got {per_teacher_lens}" + ) + self.num_teachers = len(self.projection_matrix_paths) + # ------------------------------------------------------------------ + # CSD state. Each cross-tokenizer teacher's + # KD term is the prefix-bidir partition KL; same-vocab teachers keep the + # direct-KL path and hold no CSD state. Per-teacher artifact caches + # (common-index set, prefix-support index, pseudo-target tables) are + # keyed by teacher index so any number of cross-tok teachers compose. + # ------------------------------------------------------------------ + self.common_indices_from_subtoks = bool(cfg.common_indices_from_subtoks) + # Per-teacher pseudo-target table paths (forward: student->teacher + # chains; reverse: teacher->student chains). None for same-vocab + # teachers. Runtime-injected symmetrically to projection_matrix_paths. + self.student_to_teacher_table_paths: list[Optional[str]] = ( + list(cfg.student_to_teacher_table_paths) + if cfg.student_to_teacher_table_paths is not None + else [None for _ in range(self.num_teachers)] + ) + self.teacher_to_student_table_paths: list[Optional[str]] = ( + list(cfg.teacher_to_student_table_paths) + if cfg.teacher_to_student_table_paths is not None + else [None for _ in range(self.num_teachers)] + ) + if ( + len(self.student_to_teacher_table_paths) != self.num_teachers + or len(self.teacher_to_student_table_paths) != self.num_teachers + ): + raise ValueError("CSD table paths must contain one entry per teacher") + # Optional per-microbatch loss dump for parity comparison + # (NRL_CSD_LOSS_DUMP_DIR). Raw floats from the loss-compute site. + self._loss_dump_dir = os.environ.get("NRL_CSD_LOSS_DUMP_DIR") + self._loss_dump_records: list[dict[str, Any]] = [] + self._loss_dump_call_idx = 0 + # Keep the validated loss settings alongside per-teacher metadata. + self.cfg = cfg + # Lazy CSD caches, populated on first call, keyed by (device, teacher_idx). + self._csd_common_indices_per_device: dict[ + tuple[torch.device, int], tuple[torch.Tensor, torch.Tensor] + ] = {} + self._csd_teacher_to_common_student_per_device: dict[ + tuple[torch.device, int], torch.Tensor + ] = {} + self._csd_prefix_index_per_device: dict[ + tuple[torch.device, int], dict[str, Any] + ] = {} + # Per-teacher pseudo-target tables, loaded lazily. idx -> (subtoks, + # lengths); forward = student->teacher chains, reverse = teacher->student. + self._csd_fwd_subtoks_per_teacher: dict[ + int, tuple[torch.Tensor, torch.Tensor] + ] = {} + self._csd_rev_subtoks_per_teacher: dict[ + int, tuple[torch.Tensor, torch.Tensor] + ] = {} + # The materialized projection matrix and the derived exact-map + # partition both live in process-local caches in + # ``x_token.loss_utils`` (see ``get_sparse_projection_matrix``, + # ``get_topk_projection``, ``build_exact_token_map``), not on + # this instance. That keeps the driver-side ``loss_fn`` free of + # any large CUDA tensors and lets multiple loss instances on + # the same worker share one load. + + def _teacher_is_same_vocab(self, i: int) -> bool: + """Same-tokenizer teachers supply neither CSD tables nor a projection.""" + return ( + self.projection_matrix_paths[i] is None + and self.student_to_teacher_table_paths[i] is None + ) + + @staticmethod + def _student_kd_mask(align: LocalizedAlignment) -> torch.Tensor: + """Return the semantic KD mask, with legacy text-batch fallback.""" + mask = align.student_kd_token_mask + if mask is None: + mask = align.student_token_mask + if mask is None: + raise ValueError("localized alignment is missing a student KD token mask") + return to_local_if_dtensor(mask) + + def __call__( + self, + data: BatchedDataDict[CSDLossDataDict], + global_valid_seqs: torch.Tensor, + global_valid_toks: torch.Tensor, + logits: torch.Tensor, + student_logits_contig: Optional[torch.Tensor], + teacher_full_logits_by_idx: dict[int, torch.Tensor], + aligns_by_idx: dict[int, LocalizedAlignment], + *, + student_next_token_logprobs: Optional[torch.Tensor] = None, + student_next_token_mask: Optional[torch.Tensor] = None, + teacher_sparse_logits_by_idx: Optional[dict[int, SparseTeacherLogits]] = None, + teacher_row_payloads_by_idx: Optional[dict[int, TeacherRowPayload]] = None, + dense_reconstruction_fallbacks_by_idx: Optional[dict[int, int]] = None, + global_valid_chunks_by_idx: Optional[dict[int, torch.Tensor]] = None, + global_valid_kd_toks: Optional[torch.Tensor] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + dp_cp_group: Optional[torch.distributed.ProcessGroup] = None, + cp_normalized_by_caller: bool = False, + ) -> tuple[torch.Tensor, dict[str, Any]]: + """Compute the (multi-teacher) cross-tokenizer distillation loss. + + Per-teacher KD terms are aggregated per ``kd_loss_mode`` and combined + with a single student next-token CE term — dynamic-scaled when + ``dynamic_loss_scaling`` is set (KD rescaled to match the detached CE + magnitude, ``kl_loss_weight`` / ``ce_loss_scale`` ignored), else + fixed-weighted. The single-teacher path is just ``num_teachers == 1``. + + ``student_logits_contig`` (CP-relaid), the per-teacher ``aligns_by_idx``, + and the dense and sparse teacher-logit dictionaries are precomputed in + ``prepare_loss_input``. The Automodel CP path also supplies its + sequence-local CE inputs; Megatron keeps the raw TP/CP-sharded + ``logits`` for its memory-efficient CE path. Production callers provide + ``global_valid_chunks_by_idx`` from the full batch before microbatching. + + ``dp_cp_group`` is the group the global normalizers reduce over. It is + stashed on the instance rather than threaded through the four levels of + per-teacher dispatch below; ``__call__`` is the single entry point and is + not re-entrant, so the value is always the one for the call in flight. + ``None`` keeps the legacy ``WORLD`` behaviour (single-process tests, and + the DTensor path, which has no pipeline axis to deadlock on). + + ``cp_normalized_by_caller`` is set when the caller divides the returned + loss by the CP size (the replicated Megatron path). A same-tokenizer KD + term covers only this rank's contiguous window, so it is then summed + across CP (with a SUM backward) to keep its full weight after the + division. + + With ``owner_local`` (eligible Megatron students), cross-tokenizer + teachers arrive as ``teacher_row_payloads_by_idx`` and ``logits`` stays + in its CP shard. CE and every KD term are then CP-partial: each CP rank + returns its own share, the caller must not divide by the CP size, and + dynamic scaling and the reported metrics use detached CP sums. + """ + self._dp_cp_group = dp_cp_group + self._sum_window_kd_over_cp = cp_normalized_by_caller + teacher_sparse_logits_by_idx = teacher_sparse_logits_by_idx or {} + dense_reconstruction_fallbacks_by_idx = ( + dense_reconstruction_fallbacks_by_idx or {} + ) + kd_normalizer = ( + global_valid_toks if global_valid_kd_toks is None else global_valid_kd_toks + ) + if self.owner_local: + ce_loss = self._compute_owner_local_ce( + logits, data, global_valid_toks, tp_group=tp_group, cp_group=cp_group + ) + total_kd, per_teacher_metrics = self._owner_local_sum_kd( + logits, + student_logits_contig, + teacher_full_logits_by_idx, + teacher_row_payloads_by_idx or {}, + aligns_by_idx, + kd_normalizer, + global_valid_chunks_by_idx, + tp_group=tp_group, + cp_group=cp_group, + ) + reported_kd = self._detached_cp_sum(total_kd, cp_group) + reported_ce = self._detached_cp_sum(ce_loss, cp_group) + loss, reported_loss, kl_scale = self._combine_kd_and_ce( + total_kd, ce_loss, reported_kd, reported_ce + ) + accuracy = self._owner_local_next_token_accuracy( + logits, data, tp_group=tp_group, cp_group=cp_group + ) + metrics = self._loss_metrics( + data, + loss=reported_loss, + kd=reported_kd, + ce=reported_ce, + kl_scale=kl_scale, + accuracy=accuracy, + dense_reconstruction_fallbacks_by_idx=dense_reconstruction_fallbacks_by_idx, + ) + metrics.update(per_teacher_metrics) + metrics.update(self._teacher_routing_metrics(data)) + return loss, metrics + + if student_logits_contig is None: + raise ValueError("Replicated CSD evaluation needs the student window.") + ce_loss = self._compute_ce( + logits, + data, + global_valid_toks, + student_next_token_logprobs=student_next_token_logprobs, + student_next_token_mask=student_next_token_mask, + student_logits_contig=student_logits_contig, + tp_group=tp_group, + cp_group=cp_group, + ) + + if self.kd_loss_mode == "sum": + total_kd, per_teacher_metrics = self._sum_kd( + student_logits_contig, + data, + teacher_full_logits_by_idx, + aligns_by_idx, + kd_normalizer, + teacher_sparse_logits_by_idx=teacher_sparse_logits_by_idx, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, + tp_group=tp_group, + cp_group=cp_group, + ) + elif self.kd_loss_mode == "averaged_logits": + total_kd, per_teacher_metrics = self._averaged_logits_kd( + student_logits_contig, + data, + teacher_full_logits_by_idx, + aligns_by_idx, + kd_normalizer, + teacher_sparse_logits_by_idx=teacher_sparse_logits_by_idx, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, + tp_group=tp_group, + cp_group=cp_group, + ) + elif self.kd_loss_mode == "select_teacher": + total_kd, per_teacher_metrics = self._select_teacher_kd( + student_logits_contig, + data, + teacher_full_logits_by_idx, + aligns_by_idx, + kd_normalizer, + teacher_sparse_logits_by_idx=teacher_sparse_logits_by_idx, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, + tp_group=tp_group, + cp_group=cp_group, + ) + else: + raise ValueError(f"Unknown kd_loss_mode: {self.kd_loss_mode!r}") + + # Combine the aggregated KD term with the single student CE term. + loss, _, kl_scale = self._combine_kd_and_ce( + total_kd, ce_loss, total_kd.detach(), ce_loss.detach() + ) + + # Next-token accuracy on the student side (quick per-step signal), masked + # to valid tokens. Computed once on the student from the shared CP-relaid + # fields (carried on every teacher's align); the CP-aware shift pairs + # predictors with the right labels under load-balanced sharding. + align0 = aligns_by_idx[0] + assert align0.student_input_ids is not None + assert align0.student_token_mask is not None + accuracy = next_token_accuracy( + student_logits_contig, + input_ids=align0.student_input_ids, + token_mask=align0.student_token_mask, + sample_mask=data["sample_mask"], + tp_group=tp_group, + cp_group=cp_group, + ) + + metrics = self._loss_metrics( + data, + loss=loss, + kd=total_kd, + ce=ce_loss, + kl_scale=kl_scale, + accuracy=accuracy, + dense_reconstruction_fallbacks_by_idx=dense_reconstruction_fallbacks_by_idx, + ) + metrics.update(per_teacher_metrics) + metrics.update(self._teacher_routing_metrics(data)) + return loss, metrics + + def _combine_kd_and_ce( + self, + total_kd: torch.Tensor, + ce_loss: torch.Tensor, + reported_kd: torch.Tensor, + reported_ce: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return ``(loss, reported_loss, kl_scale)``. + + ``reported_kd`` / ``reported_ce`` are the detached complete-batch terms + (CP sums when the terms are CP-partial); the dynamic KD scale is taken + from them so every CP rank applies the same ratio. + """ + if self.dynamic_loss_scaling: + # loss = sg(ce/kd) * kd + ce; user kl_loss_weight / ce_loss_scale + # are intentionally ignored in this branch. + kd_detached = reported_kd.abs() + ce_detached = reported_ce.abs() + kl_scale = torch.where( + kd_detached > 0, + ce_detached / kd_detached, + torch.ones_like(kd_detached), + ) + return ( + kl_scale * total_kd + ce_loss, + kl_scale * reported_kd + reported_ce, + kl_scale, + ) + kl_scale = torch.tensor(1.0, device=total_kd.device, dtype=total_kd.dtype) + return ( + self.kl_loss_weight * total_kd + self.ce_loss_scale * ce_loss, + self.kl_loss_weight * reported_kd + self.ce_loss_scale * reported_ce, + kl_scale, + ) + + @staticmethod + def _loss_metrics( + data: BatchedDataDict[CSDLossDataDict], + *, + loss: torch.Tensor, + kd: torch.Tensor, + ce: torch.Tensor, + kl_scale: torch.Tensor, + accuracy: torch.Tensor, + dense_reconstruction_fallbacks_by_idx: dict[int, int], + ) -> dict[str, Any]: + """Aggregate loss metrics shared by the replicated and owner-local paths.""" + metrics: dict[str, Any] = { + "loss": loss.item(), + # Aggregate KD term (kept under ``kl_loss`` so existing trainer + # metric handling continues to work); per-teacher terms are suffixed + # ``_t{i}``. + "kl_loss": kd.item(), + "ce_loss": ce.item(), + "kl_loss_scale": kl_scale.item(), + "accuracy": accuracy.item(), + "num_valid_samples": to_local_if_dtensor(data["sample_mask"]).sum().item(), + "ipc_reconstruction_fallbacks": sum( + dense_reconstruction_fallbacks_by_idx.values() + ), + } + for teacher_idx, fallback_count in sorted( + dense_reconstruction_fallbacks_by_idx.items() + ): + metrics[f"ipc_reconstruction_fallbacks_t{teacher_idx}"] = fallback_count + return metrics + + def _teacher_routing_metrics( + self, + data: BatchedDataDict[CSDLossDataDict], + ) -> dict[str, int]: + """Count the logical samples and input tokens routed to each teacher. + + Every teacher forward receives the microbatch, regardless of how its KD + term is aggregated later. Cross-tokenizer teachers are counted on their + own tokenization; same-tokenizer teachers reuse the student tokenization. + Masked padding rows and their tokens do not contribute. + """ + sample_mask = to_local_if_dtensor(data["sample_mask"]).bool() + routed_samples = int(sample_mask.sum().item()) + metrics: dict[str, int] = {} + for teacher_idx in range(self.num_teachers): + token_mask_key = ( + "token_mask" + if self._teacher_is_same_vocab(teacher_idx) + else f"teacher_{teacher_idx}_token_mask" + ) + token_mask = to_local_if_dtensor(data[token_mask_key]).bool() + routed_tokens = int((token_mask & sample_mask.unsqueeze(-1)).sum().item()) + metrics[f"teacher_{teacher_idx}/routed_samples"] = routed_samples + metrics[f"teacher_{teacher_idx}/routed_tokens"] = routed_tokens + return metrics + + # ------------------------------------------------------------------ # + # Multi-teacher aggregation + # ------------------------------------------------------------------ # + def _compute_teacher_kd( + self, + i: int, + student_logits_contig: torch.Tensor, + data: BatchedDataDict[CSDLossDataDict], + teacher_full_logits_by_idx: dict[int, torch.Tensor], + aligns_by_idx: dict[int, LocalizedAlignment], + global_valid_toks: torch.Tensor, + *, + teacher_sparse_logits_by_idx: dict[int, SparseTeacherLogits], + tp_group: Optional[torch.distributed.ProcessGroup], + cp_group: Optional[torch.distributed.ProcessGroup], + global_valid_chunks_by_idx: Optional[dict[int, torch.Tensor]] = None, + ) -> tuple[torch.Tensor, dict[str, Any]]: + """KD term for teacher ``i`` plus its (unsuffixed) metrics. + + Dispatches on tokenizer kind: same-vocab -> direct top-k per-position KL; + cross-tokenizer -> the CSD prefix-bidir partition KL over teacher ``i``'s + localized alignment. Both consume the shared CP-relaid student logits and + route TP/CP through the parameterized loss-mode helpers. + """ + has_full_logits = i in teacher_full_logits_by_idx + has_sparse_logits = i in teacher_sparse_logits_by_idx + if has_full_logits == has_sparse_logits: + raise ValueError( + f"Teacher {i} must provide exactly one dense or sparse logits " + f"payload; dense={has_full_logits}, sparse={has_sparse_logits}." + ) + if self._teacher_is_same_vocab(i): + if has_sparse_logits: + raise ValueError( + f"Same-vocab teacher {i} requires full teacher logits." + ) + return self._compute_same_vocab_kl( + i, + student_logits_contig, + teacher_full_logits_by_idx[i], + aligns_by_idx[i], + global_valid_toks, + tp_group=tp_group, + cp_group=cp_group, + ) + + global_valid_chunks = ( + global_valid_chunks_by_idx[i] + if global_valid_chunks_by_idx is not None + else None + ) + + kd, csd_metrics = self._compute_csd_loss( + i, + student_logits_contig, + teacher_full_logits_by_idx.get(i), + aligns_by_idx[i], + teacher_sparse_payload=teacher_sparse_logits_by_idx.get(i), + teacher_vocab_size=self.teacher_vocab_sizes[i], + tp_group=tp_group, + cp_group=cp_group, + global_valid_chunks=global_valid_chunks, + ) + # Surface the KD value under the shared per-teacher metric key and keep + # the CSD chunk diagnostics; the dispatcher-level loss keys are dropped + # (the aggregate ``loss`` / ``kl_loss`` are set by ``__call__``). + metrics: dict[str, Any] = {"kl_loss": kd.item()} + for key in ( + "kl_common_per_chunk", + "kl_partition_last_per_chunk", + "top1_acc_per_chunk", + "num_common_chunks", + "num_mismatch_chunks", + ): + if key in csd_metrics: + metrics[key] = csd_metrics[key] + return kd, metrics + + def _compute_same_vocab_kl( + self, + i: int, + student_logits_contig: torch.Tensor, + teacher_full_logits: torch.Tensor, + align: LocalizedAlignment, + global_valid_toks: torch.Tensor, + *, + tp_group: Optional[torch.distributed.ProcessGroup], + cp_group: Optional[torch.distributed.ProcessGroup], + ) -> tuple[torch.Tensor, dict[str, Any]]: + """Direct top-k per-position KL for a same-tokenizer teacher. + + Identical tokenizer => teacher tokens == student tokens (identity + position alignment), so no projection / no chunk-averaging. The top-k + vocab list is chosen from KD positions only (the semantic KD mask). The + reduction is the KD-position mean normalized by ``global_valid_toks`` + (the KD-token count), scaled by ``T**2``. + """ + kd = same_tokenizer_topk_kl( + student_logits_contig, + teacher_full_logits, + kd_mask=self._student_kd_mask(align), + sample_mask=align.sample_mask, + num_kd_tokens=global_valid_toks, + vocab_size=self.student_vocab_size, + vocab_topk=self.vocab_topk, + temperature=self.temperature, + reverse_kl=self.reverse_kl, + tp_group=tp_group, + cp_group=cp_group, + ) + if self._sum_window_kd_over_cp: + kd = group_all_reduce_sum_with_grad_backward_sum(kd, cp_group) + return kd, {"kl_loss": kd.item()} + + def _sum_kd( + self, + student_logits_contig: torch.Tensor, + data: BatchedDataDict[CSDLossDataDict], + teacher_full_logits_by_idx: dict[int, torch.Tensor], + aligns_by_idx: dict[int, LocalizedAlignment], + global_valid_toks: torch.Tensor, + *, + teacher_sparse_logits_by_idx: dict[int, SparseTeacherLogits], + tp_group: Optional[torch.distributed.ProcessGroup], + cp_group: Optional[torch.distributed.ProcessGroup], + global_valid_chunks_by_idx: Optional[dict[int, torch.Tensor]] = None, + ) -> tuple[torch.Tensor, dict[str, Any]]: + """Weighted sum: ``total_kd = Σ_i weight_i · KD_i``. + + Weights are static (config ``weight``) or dynamic (``sum_weights_metric``). + When ``normalize_teacher_by_vocab`` is set, each teacher's KD is + additionally scaled by ``log(V_t_i) / log(min_j V_t_j)``. + """ + device = student_logits_contig.device + if self.sum_weights_metric is not None: + if teacher_sparse_logits_by_idx: + raise ValueError( + "sum_weights_metric requires full teacher logits and cannot " + "be used with sparse teacher IPC." + ) + weights = self._compute_dynamic_weights( + data, teacher_full_logits_by_idx, aligns_by_idx + ) + else: + weights = [ + torch.tensor( + self.teacher_weights[i], + device=device, + dtype=student_logits_contig.dtype, + ) + for i in range(self.num_teachers) + ] + + if self.normalize_teacher_by_vocab: + temp_weight = torch.log( + torch.tensor(float(min(self.teacher_vocab_sizes)), device=device) + ) + + total_kd: Optional[torch.Tensor] = None + per_metrics: dict[str, Any] = {} + # Deterministic teacher order: each teacher's KD fires its own + # collectives, so the order must match across ranks. + for i in range(self.num_teachers): + kd_i, m_i = self._compute_teacher_kd( + i, + student_logits_contig, + data, + teacher_full_logits_by_idx, + aligns_by_idx, + global_valid_toks, + teacher_sparse_logits_by_idx=teacher_sparse_logits_by_idx, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, + tp_group=tp_group, + cp_group=cp_group, + ) + weighted = kd_i * weights[i] + if self.normalize_teacher_by_vocab: + v_scale = ( + torch.log( + torch.tensor(float(self.teacher_vocab_sizes[i]), device=device) + ) + / temp_weight + ) + weighted = weighted * v_scale + total_kd = weighted if total_kd is None else total_kd + weighted + per_metrics[f"teacher_{i}/weighted_kl"] = float(weighted.item()) + for k, v in m_i.items(): + per_metrics[f"{k}_t{i}"] = v + per_metrics[f"weight_t{i}"] = float(weights[i].item()) + assert total_kd is not None + return total_kd, per_metrics + + def _averaged_logits_kd( + self, + student_logits_contig: torch.Tensor, + data: BatchedDataDict[CSDLossDataDict], + teacher_full_logits_by_idx: dict[int, torch.Tensor], + aligns_by_idx: dict[int, LocalizedAlignment], + global_valid_toks: torch.Tensor, + *, + teacher_sparse_logits_by_idx: dict[int, SparseTeacherLogits], + tp_group: Optional[torch.distributed.ProcessGroup], + cp_group: Optional[torch.distributed.ProcessGroup], + global_valid_chunks_by_idx: Optional[dict[int, torch.Tensor]] = None, + ) -> tuple[torch.Tensor, dict[str, Any]]: + """Convex-weighted average of teacher logits, then one direct KL. + + Valid only when all teachers are same-tokenizer (no projection) and ship + full logits of identical shape. Otherwise falls back to a plain + static-weight sum (no dynamic weights, no ``normalize_teacher_by_vocab``). + """ + full = [teacher_full_logits_by_idx.get(i) for i in range(self.num_teachers)] + # Direct per-position KL is only valid when every teacher shares the + # student's tokenizer (no projection matrix) *and* ships full logits of + # identical shape. Two cross-tokenizer teachers can have matching shapes + # yet still need the projection/alignment path, so the shape check alone + # is insufficient. + same_tokenizer = all( + self._teacher_is_same_vocab(i) for i in range(self.num_teachers) + ) + same_shape = all(f is not None for f in full) and ( + len({tuple(f.shape) for f in full if f is not None}) == 1 + ) + if not (same_tokenizer and same_shape): + total_kd: Optional[torch.Tensor] = None + per_metrics: dict[str, Any] = {} + for i in range(self.num_teachers): + kd_i, m_i = self._compute_teacher_kd( + i, + student_logits_contig, + data, + teacher_full_logits_by_idx, + aligns_by_idx, + global_valid_toks, + teacher_sparse_logits_by_idx=teacher_sparse_logits_by_idx, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, + tp_group=tp_group, + cp_group=cp_group, + ) + w = self.teacher_weights[i] + weighted = kd_i * w + total_kd = weighted if total_kd is None else total_kd + weighted + per_metrics[f"teacher_{i}/weighted_kl"] = float(weighted.item()) + for k, v in m_i.items(): + per_metrics[f"{k}_t{i}"] = v + per_metrics[f"weight_t{i}"] = float(w) + assert total_kd is not None + return total_kd, per_metrics + + total_w = sum(self.teacher_weights) + avg: Optional[torch.Tensor] = None + for i, f in enumerate(full): + assert f is not None + contrib = f.float() * (self.teacher_weights[i] / total_w) + avg = contrib if avg is None else avg + contrib + assert avg is not None + align = aligns_by_idx[0] + kd = same_tokenizer_full_vocab_kl( + student_logits_contig, + avg, + kd_mask=self._student_kd_mask(align), + sample_mask=align.sample_mask, + num_kd_tokens=global_valid_toks, + vocab_size=self.student_vocab_size, + temperature=self.temperature, + reverse_kl=self.reverse_kl, + tp_group=tp_group, + cp_group=cp_group, + ) + if self._sum_window_kd_over_cp: + kd = group_all_reduce_sum_with_grad_backward_sum(kd, cp_group) + per_metrics: dict[str, Any] = {"kl_loss": kd.item()} + for i, weight in enumerate(self.teacher_weights): + per_metrics[f"teacher_{i}/weighted_kl"] = float( + kd.item() * weight / total_w + ) + return kd, per_metrics + + def _dp_global_masked_mean( + self, values: torch.Tensor, mask: torch.Tensor + ) -> torch.Tensor: + """Masked mean of ``values`` over the *process-global* valid count. + + Teacher selection / dynamic weighting must be identical on every rank: a + rank-local mean lets ranks pick a different teacher / different weights, + and the per-teacher KD's collectives then see divergent participation + (deadlock when one rank's choice fires a collective another's does not). + All-reduce the masked sum and the mask count over the full group so every + rank gets the same score. The result is detached (it gates selection / + weighting and is not back-propagated). + + Reduces over ``self._dp_cp_group`` when set. That group is confined to + one pipeline stage; ``WORLD`` would deadlock under PP, where only the + last stage runs the loss and reaches this collective at all. + """ + group = ( + self._dp_cp_group + if self._dp_cp_group is not None + else torch.distributed.group.WORLD + ) + num = group_all_reduce_sum((values * mask).sum(), group=group) + den = group_all_reduce_sum(mask.sum(), group=group).clamp(min=1.0) + return num / den + + def _select_teacher_kd( + self, + student_logits_contig: torch.Tensor, + data: BatchedDataDict[CSDLossDataDict], + teacher_full_logits_by_idx: dict[int, torch.Tensor], + aligns_by_idx: dict[int, LocalizedAlignment], + global_valid_toks: torch.Tensor, + *, + teacher_sparse_logits_by_idx: dict[int, SparseTeacherLogits], + tp_group: Optional[torch.distributed.ProcessGroup], + cp_group: Optional[torch.distributed.ProcessGroup], + global_valid_chunks_by_idx: Optional[dict[int, torch.Tensor]] = None, + ) -> tuple[torch.Tensor, dict[str, Any]]: + """Use only the teacher with the lowest next-token CE on its own tokens.""" + if teacher_sparse_logits_by_idx: + raise ValueError( + "kd_loss_mode='select_teacher' requires full teacher logits and " + "cannot be used with sparse teacher IPC." + ) + with torch.no_grad(): + ces: list[float] = [] + for i in range(self.num_teachers): + t_logits, t_ids, t_mask = self._teacher_score_inputs( + i, data, teacher_full_logits_by_idx, aligns_by_idx + ) + ce_pos = torch.nn.functional.cross_entropy( + t_logits[:, :-1].reshape(-1, t_logits.shape[-1]).float(), + t_ids[:, 1:].reshape(-1), + reduction="none", + ) + mask = ( + t_mask[:, 1:].float() * data["sample_mask"].unsqueeze(-1).float() + ).reshape(-1) + ces.append(self._dp_global_masked_mean(ce_pos, mask).item()) + best = int(min(range(self.num_teachers), key=lambda j: ces[j])) + + kd, m = self._compute_teacher_kd( + best, + student_logits_contig, + data, + teacher_full_logits_by_idx, + aligns_by_idx, + global_valid_toks, + teacher_sparse_logits_by_idx=teacher_sparse_logits_by_idx, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, + tp_group=tp_group, + cp_group=cp_group, + ) + per_metrics: dict[str, Any] = {f"{k}_t{best}": v for k, v in m.items()} + for i in range(self.num_teachers): + per_metrics[f"teacher_{i}/weighted_kl"] = 0.0 + per_metrics[f"teacher_{best}/weighted_kl"] = float(kd.item()) + per_metrics["selected_teacher"] = best + return kd, per_metrics + + def _teacher_score_inputs( + self, + i: int, + data: BatchedDataDict[CSDLossDataDict], + teacher_full_logits_by_idx: dict[int, torch.Tensor], + aligns_by_idx: dict[int, LocalizedAlignment], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return ``(logits, input_ids, token_mask)`` for teacher ``i``'s CE / weight-metric scores. + + The token mask is over the tokenization the score is computed on: the + shared student tokens (CP-relaid) for a same-vocab teacher, teacher ``i``'s + own otherwise. Every teacher ships full logits, so the full distribution + is always available. + """ + if self._teacher_is_same_vocab(i): + align = aligns_by_idx[i] + assert align.student_input_ids is not None + ids = to_local_if_dtensor(align.student_input_ids) + token_mask = self._student_kd_mask(align) + else: + ids = to_local_if_dtensor(data[f"teacher_{i}_input_ids"]) + token_mask = to_local_if_dtensor(data[f"teacher_{i}_token_mask"]) + return teacher_full_logits_by_idx[i], ids, token_mask + + def _compute_dynamic_weights( + self, + data: BatchedDataDict[CSDLossDataDict], + teacher_full_logits_by_idx: dict[int, torch.Tensor], + aligns_by_idx: dict[int, LocalizedAlignment], + ) -> list[torch.Tensor]: + """Sequence-level dynamic teacher weights via ``sum_weights_metric``. + + Per teacher computes a scalar score (``ce`` -> -CE, ``entropy`` -> + -entropy, ``max_prob`` -> max prob; higher = more trusted), optionally + rescaled by ``log(V_t_i)/log(min_j V_t_j)``, then ``softmax(alpha * + scores)`` across teachers. + """ + device = data["input_ids"].device + if self.normalize_teacher_by_vocab: + temp_weight = torch.log( + torch.tensor(float(min(self.teacher_vocab_sizes)), device=device) + ) + scores: list[torch.Tensor] = [] + for i in range(self.num_teachers): + t_logits, t_ids, t_mask = self._teacher_score_inputs( + i, data, teacher_full_logits_by_idx, aligns_by_idx + ) + score = self._teacher_weight_score( + t_logits, t_ids, t_mask, data["sample_mask"] + ) + if self.normalize_teacher_by_vocab: + v_log = torch.log( + torch.tensor(float(self.teacher_vocab_sizes[i]), device=device) + ) + score = score * (v_log / temp_weight) + scores.append(score) + weights = torch.softmax(self.alpha * torch.stack(scores), dim=0) + return [weights[i] for i in range(self.num_teachers)] + + def _teacher_weight_score( + self, + t_logits: torch.Tensor, + t_ids: torch.Tensor, + t_mask: torch.Tensor, + sample_mask: torch.Tensor, + ) -> torch.Tensor: + """Scalar weight-metric score for one teacher (higher = more trusted). + + Padded positions and masked-out samples are excluded, so long-padded + batches don't let near-uniform padding logits dominate the score. + """ + samp = to_local_if_dtensor(sample_mask).unsqueeze(-1).float() + if self.sum_weights_metric == "ce": + ce_pos = torch.nn.functional.cross_entropy( + t_logits[:, :-1].reshape(-1, t_logits.shape[-1]).float(), + t_ids[:, 1:].reshape(-1), + reduction="none", + ) + mask = (t_mask[:, 1:].float() * samp).reshape(-1) + return -self._dp_global_masked_mean(ce_pos, mask) + mask = t_mask.float() * samp + if self.sum_weights_metric == "entropy": + probs = torch.softmax(t_logits.float(), dim=-1) + entropy = -(probs * torch.log(probs + 1e-10)).sum(dim=-1) + return -self._dp_global_masked_mean(entropy, mask) + if self.sum_weights_metric == "max_prob": + probs = torch.softmax(t_logits.float(), dim=-1) + return self._dp_global_masked_mean(probs.max(dim=-1).values, mask) + raise ValueError(f"Unknown sum_weights_metric: {self.sum_weights_metric!r}") + + # ------------------------------------------------------------------ # + # Loss-mode implementations + # ------------------------------------------------------------------ # + def _compute_ce( + self, + logits: torch.Tensor, + data: BatchedDataDict[CSDLossDataDict], + global_valid_toks: torch.Tensor, + *, + student_next_token_logprobs: Optional[torch.Tensor] = None, + student_next_token_mask: Optional[torch.Tensor] = None, + student_logits_contig: Optional[torch.Tensor] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + ) -> torch.Tensor: + """Next-token CE on the student side (TP/CP handled by the helpers). + + The Automodel CP path supplies sequence-local target log-probabilities, + avoiding full-tensor gathers. DTensor logits otherwise carry their own + device mesh, so the vocab-parallel log-prob helper resolves TP/CP from + it. Megatron returns a *plain* tensor that is vocab-sharded (TP) and + load-balanced sequence-sharded (CP). Compute only the target + log-probabilities in fixed sequence chunks and gather the resulting + ``[B, S]`` scalars across CP; never materialize a rank-3 full-sequence or + full-vocabulary tensor. Ordinary unsharded tensors keep their existing + path. + """ + if student_next_token_logprobs is not None: + assert student_next_token_mask is not None + label_mask = student_next_token_mask.to( + student_next_token_logprobs.dtype + ) * to_local_if_dtensor(data["sample_mask"]).to( + student_next_token_logprobs.device + ).unsqueeze(-1) + local_ce = masked_mean( + -student_next_token_logprobs, + label_mask, + global_normalization_factor=global_valid_toks, + ) + # Forward SUM gives every CP rank the same full-sequence CE (and + # therefore the same dynamic KD scale). Identity backward keeps the + # disjoint rank-local windows at a single gradient fanout. + return group_all_reduce_sum_with_grad(local_ce, cp_group) + + tp_size = ( + torch.distributed.get_world_size(tp_group) + if tp_group is not None and torch.distributed.is_initialized() + else 1 + ) + cp_size = ( + torch.distributed.get_world_size(cp_group) + if cp_group is not None and torch.distributed.is_initialized() + else 1 + ) + if ( + not isinstance(logits, DTensor) + and tp_group is not None + and (tp_size > 1 or cp_size > 1) + ): + vocab_parallel_rank = torch.distributed.get_rank(tp_group) + per_token_ce = -get_next_token_logprobs_from_logits( + input_ids=data["input_ids"], + next_token_logits=logits, + seq_index=data.get("seq_index"), + vocab_parallel_rank=vocab_parallel_rank, + vocab_parallel_group=tp_group, + context_parallel_group=cp_group, + chunk_size=CSD_STUDENT_LOSS_CHUNK_SIZE, + ) + else: + per_token_ce = student_next_token_ce( + logits, + input_ids=data["input_ids"], + seq_index=data.get("seq_index"), + ) + label_mask = ce_label_mask( + token_mask=data["token_mask"], + sample_mask=data["sample_mask"], + ce_seq_len=per_token_ce.shape[1], + dtype=per_token_ce.dtype, + ) + return masked_mean( + per_token_ce, + label_mask, + global_normalization_factor=global_valid_toks, + ) + + # CSD partition-loss helpers. + def _get_csd_teacher_to_common_student( + self, + device: torch.device, + i: int, + v_t: int, + common_student_idx_t: torch.Tensor, + common_teacher_idx_t: torch.Tensor, + ) -> torch.Tensor: + cache_key = (device, i) + cached = self._csd_teacher_to_common_student_per_device.get(cache_key) + if cached is not None: + return cached + mapping = torch.full( + (int(v_t),), + -1, + dtype=torch.long, + device=device, + ) + valid = common_teacher_idx_t < int(v_t) + mapping[common_teacher_idx_t[valid]] = common_student_idx_t[valid] + self._csd_teacher_to_common_student_per_device[cache_key] = mapping + return mapping + + @staticmethod + def _lookup_sparse_teacher_logp( + topk_logits: torch.Tensor, + topk_indices: torch.Tensor, + teacher_log_z: torch.Tensor, + b_idx: torch.Tensor, + pos: torch.Tensor, + token_ids: torch.Tensor, + temperature: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Lookup teacher log-probs for explicit token ids in sparse top-k rows. + + ``topk_indices`` must be sorted ascending along the last dimension. + ``pos`` and ``token_ids`` may be either ``[K]`` or ``[K, N]``. + Missing ids return an arbitrary log-prob value and ``found=False``; + callers must mask them into the rest bucket. + """ + squeeze = False + if pos.dim() == 1: + pos = pos.unsqueeze(-1) + token_ids = token_ids.unsqueeze(-1) + squeeze = True + + original_shape = token_ids.shape + b_flat = b_idx.unsqueeze(-1).expand_as(pos).reshape(-1) + pos_flat = pos.reshape(-1) + token_flat = token_ids.reshape(-1).to(topk_indices.dtype) + out_logp = torch.empty( + (token_flat.numel(),), + device=topk_logits.device, + dtype=topk_logits.dtype, + ) + out_found = torch.empty( + (token_flat.numel(),), + device=topk_indices.device, + dtype=torch.bool, + ) + + k_top = topk_indices.shape[-1] + row_chunk = int(os.environ.get("NRL_CSD_SPARSE_LOOKUP_ROWS", "2048")) + row_chunk = max(row_chunk, 1) + for start in range(0, token_flat.numel(), row_chunk): + end = min(start + row_chunk, token_flat.numel()) + rows_idx = topk_indices[b_flat[start:end], pos_flat[start:end]] + rows_vals = topk_logits[b_flat[start:end], pos_flat[start:end]] + targets = token_flat[start:end] + + insert = torch.searchsorted(rows_idx, targets.unsqueeze(-1)) + insert = insert.squeeze(-1) + in_bounds = insert < k_top + safe_insert = insert.clamp(max=k_top - 1) + row_arange = torch.arange( + rows_idx.shape[0], + device=rows_idx.device, + dtype=torch.long, + ) + matched_idx = rows_idx[row_arange, safe_insert] + found = in_bounds & (matched_idx == targets) + raw_logits = rows_vals[row_arange, safe_insert] + flat_log_z = teacher_log_z[b_flat[start:end], pos_flat[start:end]] + + out_logp[start:end] = raw_logits / float(temperature) - flat_log_z + out_found[start:end] = found + + logp = out_logp.reshape(original_shape) + found = out_found.reshape(original_shape) + + if squeeze: + logp = logp.squeeze(-1) + found = found.squeeze(-1) + return logp, found + + def _maybe_dump_loss(self, metrics: dict[str, Any]) -> None: + """Append per-call raw loss values to a per-rank dump file. + + Activated by ``NRL_CSD_LOSS_DUMP_DIR``. One file per rank, + rewritten on each call with the full record list. Records are raw + ``loss.item()`` values from the loss-compute site — not scaled, + aggregated, or DP-summed — for loss parity comparisons. + """ + loss_dump_dir = self._loss_dump_dir + if not loss_dump_dir: + return + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + # The P-KL path emits kl_loss/ce_loss/kl_loss_scale/num_valid_pairs; + # the gold-loss path emits kl_common/l1_uncommon/num_valid_chunks; + # the CSD path emits ce_loss_per_token/kl_loss_per_token plus the + # per-chunk diagnostics. Record everything that's present so the + # same dump file format serves all — downstream comparison scripts + # read by key. + record: dict[str, Any] = { + "call_idx": self._loss_dump_call_idx, + "loss": metrics["loss"], + } + for k in ( + "kl_loss", + "ce_loss", + "kl_loss_per_chunk", + "ce_loss_per_token", + "kl_loss_scale", + "num_valid_pairs", + "kl_common", + "kl_common_per_chunk", + "kl_partition_last_per_chunk", + "l1_uncommon", + "num_valid_chunks", + ): + if k in metrics: + record[k] = metrics[k] + self._loss_dump_records.append(record) + self._loss_dump_call_idx += 1 + os.makedirs(loss_dump_dir, exist_ok=True) + torch.save( + self._loss_dump_records, + os.path.join(loss_dump_dir, f"rank{rank}.pt"), + ) + + # ------------------------------------------------------------------ # + # Loss-mode implementations + # ------------------------------------------------------------------ # + + def _csd_fwd_table(self, i: int) -> tuple[torch.Tensor, torch.Tensor]: + """Lazily load+cache teacher ``i``'s forward subtoks table. + + Forward = student->teacher sub-token chains (``subtoks`` / ``lengths`` + keys). Small CPU tensors, shared by ``_common_indices_from_subtoks`` and + the forward branch of ``_ensure_bidir_prefix_support_index``. + """ + cached = self._csd_fwd_subtoks_per_teacher.get(i) + if cached is not None: + return cached + fwd_path = self.student_to_teacher_table_paths[i] or "" + if not fwd_path or not os.path.exists(fwd_path): + raise RuntimeError( + f"teacher {i}: forward pseudo-target table required " + f"(student_to_teacher_table_paths[{i}]); got {fwd_path!r}." + ) + fwd = torch.load(fwd_path, map_location="cpu", weights_only=False) + table = (fwd["subtoks"].long(), fwd["lengths"].long()) + self._csd_fwd_subtoks_per_teacher[i] = table + return table + + def _csd_rev_table(self, i: int) -> tuple[torch.Tensor, torch.Tensor]: + """Lazily load+cache teacher ``i``'s reverse subtoks table. + + Reverse = teacher->student sub-token chains, consumed by the reverse + branch of ``_ensure_bidir_prefix_support_index``. + """ + cached = self._csd_rev_subtoks_per_teacher.get(i) + if cached is not None: + return cached + rev_path = self.teacher_to_student_table_paths[i] or "" + if not rev_path or not os.path.exists(rev_path): + raise RuntimeError( + f"teacher {i}: reverse pseudo-target table required " + f"(teacher_to_student_table_paths[{i}]); got {rev_path!r}." + ) + rev = torch.load(rev_path, map_location="cpu", weights_only=False) + table = (rev["subtoks"].long(), rev["lengths"].long()) + self._csd_rev_subtoks_per_teacher[i] = table + return table + + def _get_common_indices_csd( + self, + device: torch.device, + i: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return cached (common_student, common_teacher) for teacher ``i``. + + The exact 1-to-1 common set is sourced either from the forward subtoks + table (``common_indices_from_subtoks=True``, the CSD-preset path, no + projection matrix needed) or from the strict exact-token map of the + teacher's projection matrix. Cached per ``(device, i)``. + """ + cache_key = (device, i) + if cache_key in self._csd_common_indices_per_device: + return self._csd_common_indices_per_device[cache_key] + if self.common_indices_from_subtoks: + common_s, common_t = self._common_indices_from_subtoks(device, i) + else: + # CSD has no gold/xtoken modifier -> strict exact map + # (``xtoken_loss=False``), matching every CSD preset. + projection_path = self.projection_matrix_paths[i] + if projection_path is None: + raise ValueError( + "CSD requires a projection matrix when " + "common_indices_from_subtoks=False" + ) + exact_map = build_exact_token_map( + projection_path, + device, + xtoken_loss=False, + teacher_vocab_size=self.teacher_vocab_sizes[i], + ) + common_s = exact_map["common_student"] + common_t = exact_map["common_teacher"] + self._csd_common_indices_per_device[cache_key] = (common_s, common_t) + return common_s, common_t + + def _common_indices_from_subtoks( + self, + device: torch.device, + i: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Derive (common_student, common_teacher) from teacher ``i``'s forward subtoks table. + + A student token is exactly "common" (1-to-1 with the teacher vocab) iff + its forward sub-token chain has length 1; the paired teacher id is + ``subtoks_t[s, 0]``. This reads the same forward table the prefix + support index uses (``student_to_teacher_table_paths[i]``, keys + ``subtoks``/``lengths``), so a common-only run needs no projection + matrix. Teacher collisions (rare for true exact matches) are deduped + keeping the lowest student id, so ``common_teacher`` is unique. Result + is sorted by student index, like the projection strict exact-token map. + """ + subtoks_t, lengths_t = self._csd_fwd_table(i) + v_teacher = int(self.teacher_vocab_sizes[i]) + + common_s = (lengths_t == 1).nonzero(as_tuple=True)[0] + common_t = subtoks_t[common_s, 0] + in_bounds = (common_t >= 0) & (common_t < v_teacher) + common_s = common_s[in_bounds] + common_t = common_t[in_bounds] + + if common_t.numel() > 0: + # Sort by student id, then keep the first (lowest-student) entry per + # teacher id so common_teacher is unique. + order = torch.argsort(common_s) + common_s, common_t = common_s[order], common_t[order] + sort_t, t_order = torch.sort(common_t, stable=True) + first_of_t = torch.ones_like(sort_t, dtype=torch.bool) + first_of_t[1:] = sort_t[1:] != sort_t[:-1] + keep = torch.zeros(common_t.shape[0], dtype=torch.bool) + keep[t_order[first_of_t]] = True + common_s, common_t = common_s[keep], common_t[keep] + re_order = torch.argsort(common_s) + common_s, common_t = common_s[re_order], common_t[re_order] + + return common_s.to(device), common_t.to(device) + + def _ensure_bidir_prefix_support_index( + self, + device: torch.device, + i: int, + ) -> dict[str, Any]: + """Lazy-build teacher ``i``'s bidirectional prefix support index. + + The cache maps ``(length, prefix_tuple) -> tuple[(s_id, t_id), ...]`` + for both the forward (``student_to_teacher_table_paths[i]``) and reverse + (``teacher_to_student_table_paths[i]``) pseudo-target tables. Built once + per ``(device, i)``; the table loaders raise if a path is missing. + """ + cache_key = (device, i) + if cache_key in self._csd_prefix_index_per_device: + return self._csd_prefix_index_per_device[cache_key] + + # Small CPU tensors (shape [V, max_chain]); slices move to device on + # demand. + subtoks_t_cpu, lengths_t_cpu = self._csd_fwd_table(i) + subtoks_s_cpu, lengths_s_cpu = self._csd_rev_table(i) + + # Build the prefix-keyed dicts. + forward: dict[tuple[int, tuple[int, ...]], list[tuple[int, int]]] = ( + collections.defaultdict(list) + ) + for s_id, length in enumerate(lengths_t_cpu.tolist()): + length = int(length) + if length <= 1 or length > subtoks_t_cpu.size(1): + continue + row = subtoks_t_cpu[s_id, :length] + if bool((row < 0).any().item()): + continue + prefix = tuple(int(x) for x in row[: length - 1].tolist()) + final_t = int(row[length - 1].item()) + forward[(length, prefix)].append((int(s_id), final_t)) + + reverse: dict[tuple[int, tuple[int, ...]], list[tuple[int, int]]] = ( + collections.defaultdict(list) + ) + for t_id, length in enumerate(lengths_s_cpu.tolist()): + length = int(length) + if length <= 1 or length > subtoks_s_cpu.size(1): + continue + row = subtoks_s_cpu[t_id, :length] + if bool((row < 0).any().item()): + continue + prefix = tuple(int(x) for x in row[: length - 1].tolist()) + final_s = int(row[length - 1].item()) + reverse[(length, prefix)].append((int(t_id), final_s)) + + # Pre-dedupe. The same env override + # gate is preserved so A/B verification stays available. + _prededuped = os.environ.get("NRL_CSD_PREDEDUP_PREFIX_INDEX", "1") == "1" + + def _dedup_by_second(pairs_list): + seen_second = set() + out = [] + for p in pairs_list: + second = p[1] + if second in seen_second: + continue + seen_second.add(second) + out.append(p) + return tuple(out) + + if _prededuped: + cache: dict[str, Any] = { + "forward": { + key: _dedup_by_second(value) for key, value in forward.items() + }, + "reverse": { + key: _dedup_by_second(value) for key, value in reverse.items() + }, + "_prededuped": True, + } + else: + cache = { + "forward": {key: tuple(value) for key, value in forward.items()}, + "reverse": {key: tuple(value) for key, value in reverse.items()}, + "_prededuped": False, + } + # Cache the table sizes used by the chunk-classification skip rules. + cache["max_chain_t"] = int(subtoks_t_cpu.size(1)) + cache["max_chain_s"] = int(subtoks_s_cpu.size(1)) + self._csd_prefix_index_per_device[cache_key] = cache + return cache + + @staticmethod + def _append_rest_bucket_logp( + logp: torch.Tensor, + eps: float = 1e-12, + ) -> torch.Tensor: + """Append log probability for the complement of a support set. + + ``eps`` is the clamp floor for the complement mass. + """ + if logp.numel() == 0: + return logp + work = logp.float() if logp.dtype in (torch.float16, torch.bfloat16) else logp + support_mass = work.exp().sum(dim=-1, keepdim=True) + rest_mass = (1.0 - support_mass).clamp(min=eps) + return torch.cat([work, rest_mass.log()], dim=-1) + + @staticmethod + def _binary_power_bce_from_logp( + student_logp_with_rest: torch.Tensor, + teacher_logp_with_rest: torch.Tensor, + tau: float = 1.0, + eps: float = 1e-12, + ) -> torch.Tensor: + """Soft-label BCE on ALM's binary realized-vs-rest distribution. + + ``student_logp_with_rest`` and ``teacher_logp_with_rest`` must have + exactly two columns: realized-chain log-probability and rest-bucket + log-probability. The BCE input and target follow: + + input = p_student_realized ** (1 / tau) + target = p_teacher_realized ** (1 / tau) + + The returned tensor has the same trailing width of 2, with the realized + and rest BCE terms split so existing support/rest masking still applies. + """ + if tau <= 0.0: + raise ValueError(f"csd_alm_bce_tau must be > 0, got {tau}") + if ( + student_logp_with_rest.shape[-1] != 2 + or teacher_logp_with_rest.shape[-1] != 2 + ): + raise ValueError( + "ALM BCE expects binary realized-vs-rest support; got " + f"student shape {tuple(student_logp_with_rest.shape)} and " + f"teacher shape {tuple(teacher_logp_with_rest.shape)}" + ) + + work_dtype = ( + torch.float32 + if student_logp_with_rest.dtype in (torch.float16, torch.bfloat16) + else student_logp_with_rest.dtype + ) + eps = max(float(eps), float(torch.finfo(work_dtype).eps)) + exponent = 1.0 / float(tau) + student_real = ( + student_logp_with_rest[..., 0] + .to(work_dtype) + .exp() + .clamp(min=0.0, max=1.0) + .pow(exponent) + .clamp(min=eps, max=1.0 - eps) + ) + teacher_real = ( + teacher_logp_with_rest[..., 0] + .to(work_dtype) + .exp() + .clamp(min=0.0, max=1.0) + .pow(exponent) + .detach() + ) + realized_term = -teacher_real * student_real.log() + rest_term = -(1.0 - teacher_real) * torch.log1p(-student_real) + return torch.stack([realized_term, rest_term], dim=-1) + + @staticmethod + def _topk_mask_by_score( + candidate_mask: torch.Tensor, + scores: torch.Tensor, + k: int, + ) -> torch.Tensor: + """Keep the row-wise top-k valid candidates by score.""" + k = int(k) + if k <= 0: + return torch.zeros_like(candidate_mask, dtype=torch.bool) + candidate_mask = candidate_mask.to(torch.bool) + if candidate_mask.shape[-1] <= k: + return candidate_mask + masked_scores = torch.where( + candidate_mask, + scores, + torch.full_like(scores, -1.0e30), + ) + _, keep_idx = torch.topk(masked_scores, k=k, dim=-1) + keep_mask = torch.zeros_like(candidate_mask, dtype=torch.bool) + keep_mask.scatter_(dim=-1, index=keep_idx, value=True) + return keep_mask & candidate_mask + + @staticmethod + def _lookup_cp_sharded_teacher_logits( + teacher_logits: torch.Tensor, + batch_indices: torch.Tensor, + global_positions: torch.Tensor, + token_ids: torch.Tensor, + cp_group: Optional[torch.distributed.ProcessGroup], + ) -> torch.Tensor: + """Fetch arbitrary dense-teacher logits without gathering ``[B,T,V]``. + + ``teacher_logits`` is the caller's contiguous CP-local sequence window + with a fully reconstructed vocabulary. The query tensors broadcast to + one common shape and use global sequence positions. Under CP, exactly + one rank owns each position; that owner fills the queried values and a + no-grad SUM makes the small result available on every CP replica. This + preserves CSD's replicated numerator while avoiding a full dense + sequence gather for Qwen-sized vocabularies. + + Every caller must issue the same queries in the same order on every CP + rank. CSD constructs them only from globally gathered ids/spans and + replicated alignment metadata. + """ + batch_indices, global_positions, token_ids = torch.broadcast_tensors( + batch_indices, + global_positions, + token_ids, + ) + if batch_indices.numel() == 0: + return torch.empty_like(batch_indices, dtype=teacher_logits.dtype) + + cp_size = ( + torch.distributed.get_world_size(cp_group) + if cp_group is not None + and torch.distributed.is_initialized() + and torch.distributed.get_world_size(cp_group) > 1 + else 1 + ) + if cp_size == 1: + return teacher_logits[batch_indices, global_positions, token_ids] + + batch_size, local_seq_len, vocab_size = teacher_logits.shape + full_seq_len = local_seq_len * cp_size + # These tensors are replicated query metadata. Validate on every rank + # before entering the collective so malformed input cannot strand peers + # in all_reduce after only the owning rank raises an indexing error. + if ( + int(batch_indices.min().item()) < 0 + or int(batch_indices.max().item()) >= batch_size + ): + raise IndexError( + "dense teacher batch query is out of bounds: " + f"batch_size={batch_size}, range=" + f"[{int(batch_indices.min().item())}, " + f"{int(batch_indices.max().item())}]" + ) + if ( + int(global_positions.min().item()) < 0 + or int(global_positions.max().item()) >= full_seq_len + ): + raise IndexError( + "dense teacher sequence query is out of bounds: " + f"full_seq_len={full_seq_len}, range=" + f"[{int(global_positions.min().item())}, " + f"{int(global_positions.max().item())}]" + ) + if int(token_ids.min().item()) < 0 or int(token_ids.max().item()) >= vocab_size: + raise IndexError( + "dense teacher vocabulary query is out of bounds: " + f"vocab_size={vocab_size}, range=" + f"[{int(token_ids.min().item())}, {int(token_ids.max().item())}]" + ) + + cp_rank = torch.distributed.get_rank(cp_group) + flat_batch = batch_indices.reshape(-1) + flat_positions = global_positions.reshape(-1) + flat_tokens = token_ids.reshape(-1) + owners = torch.div(flat_positions, local_seq_len, rounding_mode="floor") + owned = owners == cp_rank + values = torch.zeros( + flat_batch.shape, + dtype=teacher_logits.dtype, + device=teacher_logits.device, + ) + with torch.no_grad(): + if bool(owned.any().item()): + values[owned] = teacher_logits[ + flat_batch[owned], + flat_positions[owned] - cp_rank * local_seq_len, + flat_tokens[owned], + ] + torch.distributed.all_reduce( + values, + op=torch.distributed.ReduceOp.SUM, + group=cp_group, + ) + return values.view(batch_indices.shape) + + @staticmethod + def _lookup_tp_cp_sharded_student_logits( + student_logits: torch.Tensor, + batch_indices: torch.Tensor, + global_positions: torch.Tensor, + token_ids: torch.Tensor, + *, + tp_group: Optional[torch.distributed.ProcessGroup], + cp_group: Optional[torch.distributed.ProcessGroup], + ) -> torch.Tensor: + """Fetch arbitrary differentiable logits from TP/CP student shards. + + ``student_logits`` has a contiguous CP-local sequence window and a + contiguous TP-local vocabulary interval. Query tensors are globally + replicated and broadcast to one shape. Exactly one TP/CP coordinate + owns each value; a TP SUM with rank-local backward and then a CP SUM + with backward SUM reproduces the former full-gather gradient semantics + under the outer ``loss / CP`` scaling, without a rank-3 gather. + """ + query_shape = torch.broadcast_shapes( + batch_indices.shape, + global_positions.shape, + token_ids.shape, + ) + if query_shape.numel() == 0: + return torch.empty( + query_shape, + dtype=student_logits.dtype, + device=student_logits.device, + ) + + tp_size = ( + torch.distributed.get_world_size(tp_group) + if tp_group is not None and torch.distributed.is_initialized() + else 1 + ) + cp_size = ( + torch.distributed.get_world_size(cp_group) + if cp_group is not None and torch.distributed.is_initialized() + else 1 + ) + batch_size, local_seq_len, local_vocab_size = student_logits.shape + full_seq_len = local_seq_len * cp_size + full_vocab_size = local_vocab_size * tp_size + + batch_min = int(batch_indices.min().item()) + batch_max = int(batch_indices.max().item()) + if batch_min < 0 or batch_max >= batch_size: + raise IndexError( + "student batch query is out of bounds: " + f"batch_size={batch_size}, range=[{batch_min}, {batch_max}]" + ) + position_min = int(global_positions.min().item()) + position_max = int(global_positions.max().item()) + if position_min < 0 or position_max >= full_seq_len: + raise IndexError( + "student sequence query is out of bounds: " + f"full_seq_len={full_seq_len}, range=" + f"[{position_min}, {position_max}]" + ) + token_min = int(token_ids.min().item()) + token_max = int(token_ids.max().item()) + if token_min < 0 or token_max >= full_vocab_size: + raise IndexError( + "student vocabulary query is out of bounds: " + f"full_vocab_size={full_vocab_size}, range=" + f"[{token_min}, {token_max}]" + ) + + if tp_size == 1 and cp_size == 1: + return student_logits[batch_indices, global_positions, token_ids] + + tp_rank = torch.distributed.get_rank(tp_group) if tp_size > 1 else 0 + cp_rank = torch.distributed.get_rank(cp_group) if cp_size > 1 else 0 + position_owners = torch.div( + global_positions, + local_seq_len, + rounding_mode="floor", + ) + token_owners = torch.div( + token_ids, + local_vocab_size, + rounding_mode="floor", + ) + local_values = student_logits[ + batch_indices, + global_positions.remainder(local_seq_len), + token_ids.remainder(local_vocab_size), + ] + # Keep the separable owner masks compact. In the dense-common path their + # shapes are [rows, 1] and [1, vocab], so broadcasting them only inside + # ``where`` avoids retaining full [rows, vocab] index/mask tensors in the + # autograd graph. ``where`` also avoids propagating a non-owner NaN via + # ``NaN * 0``. + values = torch.where(position_owners == cp_rank, local_values, 0.0) + values = torch.where(token_owners == tp_rank, values, 0.0) + values = group_all_reduce_sum_with_grad(values, tp_group) + values = group_all_reduce_sum_with_grad_backward_sum(values, cp_group) + return values + + @staticmethod + def _cp_sharded_teacher_target_in_topk( + teacher_logits: torch.Tensor, + *, + batch_index: int, + global_position: int, + token_id: int, + k: int, + cp_group: Optional[torch.distributed.ProcessGroup], + ) -> bool: + """Return exact target-in-top-k membership for a CP-sharded dense row.""" + batch_size, local_seq_len, vocab_size = teacher_logits.shape + cp_size = ( + torch.distributed.get_world_size(cp_group) + if cp_group is not None + and torch.distributed.is_initialized() + and torch.distributed.get_world_size(cp_group) > 1 + else 1 + ) + full_seq_len = local_seq_len * cp_size + if not 0 <= batch_index < batch_size: + raise IndexError( + f"teacher top-k batch index {batch_index} outside [0, {batch_size})" + ) + if not 0 <= global_position < full_seq_len: + raise IndexError( + "teacher top-k position outside the full sequence: " + f"position={global_position}, full_seq_len={full_seq_len}" + ) + if not 0 <= token_id < vocab_size: + raise IndexError( + f"teacher top-k token id {token_id} outside [0, {vocab_size})" + ) + k = min(int(k), vocab_size) + if k <= 0: + return False + if k == vocab_size: + return True + + if cp_size == 1: + topk_ids = torch.topk( + teacher_logits[batch_index, global_position], + k=k, + dim=-1, + ).indices + return bool((topk_ids == token_id).any().item()) + + cp_rank = torch.distributed.get_rank(cp_group) + owner = global_position // local_seq_len + found = torch.zeros((), dtype=torch.int32, device=teacher_logits.device) + with torch.no_grad(): + if cp_rank == owner: + local_position = global_position - owner * local_seq_len + topk_ids = torch.topk( + teacher_logits[batch_index, local_position], + k=k, + dim=-1, + ).indices + found.fill_(int(bool((topk_ids == token_id).any().item()))) + torch.distributed.all_reduce( + found, + op=torch.distributed.ReduceOp.SUM, + group=cp_group, + ) + return bool(found.item()) + + @staticmethod + def _unique_bidir_pairs_cpu( + pairs, + swap: bool = False, + assume_unique: bool = False, + ) -> tuple[list[int], list[int]]: + """CPU dedupe of ``(s_id, t_id)`` pairs. + + When ``assume_unique=True`` we skip the Python set dedupe loop and + unpack via numpy. ``swap=True`` treats input as ``(t_id, s_id)`` + (used for the reverse prefix-index branch). + """ + if not pairs: + return [], [] + if assume_unique: + arr = np.asarray(pairs, dtype=np.int64) + if swap: + return arr[:, 1].tolist(), arr[:, 0].tolist() + return arr[:, 0].tolist(), arr[:, 1].tolist() + seen_s: set[int] = set() + seen_t: set[int] = set() + s_list: list[int] = [] + t_list: list[int] = [] + for p in pairs: + if swap: + t_id = int(p[0]) + s_id = int(p[1]) + else: + s_id = int(p[0]) + t_id = int(p[1]) + if s_id in seen_s or t_id in seen_t: + continue + seen_s.add(s_id) + seen_t.add(t_id) + s_list.append(s_id) + t_list.append(t_id) + return s_list, t_list + + def _partition_kl_mismatch_batched( + self, + chunk_records, + student_logits: torch.Tensor, + teacher_logits: Optional[torch.Tensor], + student_log_z: torch.Tensor, + teacher_log_z: torch.Tensor, + temperature: float, + reverse_kl: bool, + loss_fn: str = "kl", + jsd_beta: float = 0.5, + alm_bce_tau: float = 1.0, + teacher_sparse: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + teacher_sparse_keep_realized: bool = False, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + ) -> tuple[torch.Tensor, int, int, list[int], list[bool]]: + """Batched partition-KL over a flat list of mismatch chunk records. + + Each record is the 11-tuple + ``(b, s_last, t_last, s_prefix_pred_pos, s_prefix_label_id, + t_prefix_pred_pos, t_prefix_label_id, s_ids_list, t_ids_list, + realized_s_last_id, realized_t_last_id)``. + Returns ``(loss_sum, used_count, matches, support_sizes, valid_flags)``. + """ + device = student_logits.device + accum_dtype = ( + torch.float32 + if student_logits.dtype in (torch.float16, torch.bfloat16) + else student_logits.dtype + ) + if not chunk_records: + zero = torch.tensor(0.0, device=device, dtype=accum_dtype) + return zero, 0, 0, [], [] + + K = len(chunk_records) + Pmax_s = max(len(r[3]) for r in chunk_records) + Pmax_t = max(len(r[5]) for r in chunk_records) + Smax = max(len(r[7]) for r in chunk_records) + # Sentinel that survives kl_div: exp(-1e30) underflows to 0 in fp32 so + # the masked entry contributes 0 to support-mass and to the per-row KL + # sum, while target - input stays finite (no 0*nan). + NEG_INF = -1.0e30 + + # The NumPy fast path is enabled by default. The Python-list branch + # can be selected for parity verification through the environment. + _use_np = os.environ.get("NRL_CSD_BATCHED_HELPER_NUMPY", "1") == "1" + + b_idx_list = [r[0] for r in chunk_records] + s_last_list = [r[1] for r in chunk_records] + t_last_list = [r[2] for r in chunk_records] + s_realized_list = [r[9] if len(r) > 9 else -1 for r in chunk_records] + t_realized_list = [r[10] if len(r) > 10 else -1 for r in chunk_records] + support_sizes = [len(r[7]) for r in chunk_records] + + np_accum_dtype = np.float32 if accum_dtype == torch.float32 else np.float64 + + s_pred_pos = None + s_label_id = None + s_pmask = None + t_pred_pos = None + t_label_id = None + t_pmask = None + + if _use_np: + s_ids_np = np.zeros((K, Smax), dtype=np.int64) + t_ids_np = np.zeros((K, Smax), dtype=np.int64) + smask_np = np.zeros((K, Smax), dtype=np_accum_dtype) + if Pmax_s > 0: + s_pred_pos_np = np.zeros((K, Pmax_s), dtype=np.int64) + s_label_id_np = np.zeros((K, Pmax_s), dtype=np.int64) + s_pmask_np = np.zeros((K, Pmax_s), dtype=np_accum_dtype) + if Pmax_t > 0: + t_pred_pos_np = np.zeros((K, Pmax_t), dtype=np.int64) + t_label_id_np = np.zeros((K, Pmax_t), dtype=np.int64) + t_pmask_np = np.zeros((K, Pmax_t), dtype=np_accum_dtype) + for i, r in enumerate(chunk_records): + s_pp, s_lid = r[3], r[4] + t_pp, t_lid = r[5], r[6] + s_list, t_list = r[7], r[8] + ns_p = len(s_pp) + if ns_p > 0: + s_pred_pos_np[i, :ns_p] = s_pp + s_label_id_np[i, :ns_p] = s_lid + s_pmask_np[i, :ns_p] = 1.0 + nt_p = len(t_pp) + if nt_p > 0: + t_pred_pos_np[i, :nt_p] = t_pp + t_label_id_np[i, :nt_p] = t_lid + t_pmask_np[i, :nt_p] = 1.0 + n_sup = len(s_list) + if n_sup > 0: + s_ids_np[i, :n_sup] = s_list + t_ids_np[i, :n_sup] = t_list + smask_np[i, :n_sup] = 1.0 + + b_idx = torch.from_numpy(np.asarray(b_idx_list, dtype=np.int64)).to( + device, non_blocking=True + ) + s_last = torch.from_numpy(np.asarray(s_last_list, dtype=np.int64)).to( + device, non_blocking=True + ) + t_last = torch.from_numpy(np.asarray(t_last_list, dtype=np.int64)).to( + device, non_blocking=True + ) + s_realized_last = torch.from_numpy( + np.asarray(s_realized_list, dtype=np.int64) + ).to(device, non_blocking=True) + t_realized_last = torch.from_numpy( + np.asarray(t_realized_list, dtype=np.int64) + ).to(device, non_blocking=True) + s_ids = torch.from_numpy(s_ids_np).to(device, non_blocking=True) + t_ids = torch.from_numpy(t_ids_np).to(device, non_blocking=True) + support_mask = torch.from_numpy(smask_np).to( + device, + dtype=accum_dtype, + non_blocking=True, + ) + if Pmax_s > 0: + s_pred_pos = torch.from_numpy(s_pred_pos_np).to( + device, + non_blocking=True, + ) + s_label_id = torch.from_numpy(s_label_id_np).to( + device, + non_blocking=True, + ) + s_pmask = torch.from_numpy(s_pmask_np).to( + device, + dtype=accum_dtype, + non_blocking=True, + ) + if Pmax_t > 0: + t_pred_pos = torch.from_numpy(t_pred_pos_np).to( + device, + non_blocking=True, + ) + t_label_id = torch.from_numpy(t_label_id_np).to( + device, + non_blocking=True, + ) + t_pmask = torch.from_numpy(t_pmask_np).to( + device, + dtype=accum_dtype, + non_blocking=True, + ) + else: + s_pred_pos_buf = [[0] * Pmax_s for _ in range(K)] + s_label_id_buf = [[0] * Pmax_s for _ in range(K)] + s_pmask_buf = [[0.0] * Pmax_s for _ in range(K)] + t_pred_pos_buf = [[0] * Pmax_t for _ in range(K)] + t_label_id_buf = [[0] * Pmax_t for _ in range(K)] + t_pmask_buf = [[0.0] * Pmax_t for _ in range(K)] + s_ids_buf = [[0] * Smax for _ in range(K)] + t_ids_buf = [[0] * Smax for _ in range(K)] + smask_buf = [[0.0] * Smax for _ in range(K)] + + for i, r in enumerate(chunk_records): + s_pp, s_lid = r[3], r[4] + t_pp, t_lid = r[5], r[6] + s_list, t_list = r[7], r[8] + for j, (p, lid) in enumerate(zip(s_pp, s_lid)): + s_pred_pos_buf[i][j] = p + s_label_id_buf[i][j] = lid + s_pmask_buf[i][j] = 1.0 + for j, (p, lid) in enumerate(zip(t_pp, t_lid)): + t_pred_pos_buf[i][j] = p + t_label_id_buf[i][j] = lid + t_pmask_buf[i][j] = 1.0 + n_sup = len(s_list) + for j in range(n_sup): + s_ids_buf[i][j] = s_list[j] + t_ids_buf[i][j] = t_list[j] + smask_buf[i][j] = 1.0 + + b_idx = torch.tensor(b_idx_list, device=device, dtype=torch.long) + s_last = torch.tensor( + s_last_list, + device=device, + dtype=torch.long, + ) + t_last = torch.tensor( + t_last_list, + device=device, + dtype=torch.long, + ) + s_realized_last = torch.tensor( + s_realized_list, + device=device, + dtype=torch.long, + ) + t_realized_last = torch.tensor( + t_realized_list, + device=device, + dtype=torch.long, + ) + s_ids = torch.tensor(s_ids_buf, device=device, dtype=torch.long) + t_ids = torch.tensor(t_ids_buf, device=device, dtype=torch.long) + support_mask = torch.tensor( + smask_buf, + device=device, + dtype=accum_dtype, + ) + if Pmax_s > 0: + s_pred_pos = torch.tensor( + s_pred_pos_buf, + device=device, + dtype=torch.long, + ) + s_label_id = torch.tensor( + s_label_id_buf, + device=device, + dtype=torch.long, + ) + s_pmask = torch.tensor( + s_pmask_buf, + device=device, + dtype=accum_dtype, + ) + if Pmax_t > 0: + t_pred_pos = torch.tensor( + t_pred_pos_buf, + device=device, + dtype=torch.long, + ) + t_label_id = torch.tensor( + t_label_id_buf, + device=device, + dtype=torch.long, + ) + t_pmask = torch.tensor( + t_pmask_buf, + device=device, + dtype=accum_dtype, + ) + + # ---- Student prefix chain (batched gather) ---- + if Pmax_s > 0: + assert s_pred_pos is not None + assert s_label_id is not None + assert s_pmask is not None + s_prefix_logits = ( + self._lookup_tp_cp_sharded_student_logits( + student_logits, + b_idx[:, None], + s_pred_pos, + s_label_id, + tp_group=tp_group, + cp_group=cp_group, + ) + / temperature + ) + s_prefix_scores = ( + s_prefix_logits.to(accum_dtype) + - student_log_z[b_idx[:, None], s_pred_pos] + ) + s_prefix_logp = (s_prefix_scores * s_pmask).sum(dim=-1) # (K,) + else: + s_prefix_logp = torch.zeros( + (K,), + device=device, + dtype=accum_dtype, + ) + + # ---- Teacher prefix chain (batched gather; detached) ---- + prefix_found = torch.ones((K,), device=device, dtype=torch.bool) + if Pmax_t > 0: + assert t_pred_pos is not None + assert t_label_id is not None + assert t_pmask is not None + if teacher_sparse is not None: + teacher_topk_logits, teacher_topk_indices = teacher_sparse + t_prefix_scores, t_prefix_found = self._lookup_sparse_teacher_logp( + teacher_topk_logits, + teacher_topk_indices, + teacher_log_z, + b_idx, + t_pred_pos, + t_label_id, + temperature, + ) + t_prefix_scores = torch.where( + t_prefix_found, + t_prefix_scores.to(accum_dtype), + torch.zeros_like(t_prefix_scores, dtype=accum_dtype), + ) + prefix_found = (t_prefix_found | (t_pmask <= 0)).all(dim=-1) + else: + assert teacher_logits is not None + t_prefix_logits = ( + self._lookup_cp_sharded_teacher_logits( + teacher_logits, + b_idx[:, None], + t_pred_pos, + t_label_id, + cp_group, + ) + / temperature + ) + t_prefix_scores = ( + t_prefix_logits.to(accum_dtype) + - teacher_log_z[b_idx[:, None], t_pred_pos] + ) + t_prefix_logp = (t_prefix_scores * t_pmask).sum(dim=-1) + else: + t_prefix_logp = torch.zeros( + (K,), + device=device, + dtype=accum_dtype, + ) + + # ---- Final-position support scores (batched gather) ---- + s_final_logits = ( + self._lookup_tp_cp_sharded_student_logits( + student_logits, + b_idx[:, None], + s_last[:, None], + s_ids, + tp_group=tp_group, + cp_group=cp_group, + ) + / temperature + ) + s_final_scores = s_final_logits.to(accum_dtype) - student_log_z[ + b_idx, s_last + ].unsqueeze(-1) + s_full_scores = s_final_scores + s_prefix_logp.unsqueeze(-1) + s_full_scores = torch.where( + support_mask > 0, + s_full_scores, + torch.full_like(s_full_scores, NEG_INF), + ) + + with torch.no_grad(): + if teacher_sparse is not None: + teacher_topk_logits, teacher_topk_indices = teacher_sparse + t_last_pos = t_last.unsqueeze(-1).expand_as(t_ids) + t_final_scores, t_final_found = self._lookup_sparse_teacher_logp( + teacher_topk_logits, + teacher_topk_indices, + teacher_log_z, + b_idx, + t_last_pos, + t_ids, + temperature, + ) + if teacher_sparse_keep_realized: + topk_width = int(teacher_topk_indices.shape[-1]) + alt_budget = max(topk_width - 1, 0) + realized_pair_mask = ( + (support_mask > 0) + & (s_ids == s_realized_last[:, None]) + & (t_ids == t_realized_last[:, None]) + ) + alt_candidate_mask = ( + (support_mask > 0) & t_final_found & (~realized_pair_mask) + ) + alt_keep_mask = self._topk_mask_by_score( + alt_candidate_mask, + t_final_scores.to(accum_dtype), + alt_budget, + ) + final_keep_mask = ( + realized_pair_mask & t_final_found + ) | alt_keep_mask + else: + final_keep_mask = (support_mask > 0) & t_final_found + teacher_support_mask = final_keep_mask & prefix_found.unsqueeze(-1) + t_full_scores = t_final_scores.to( + accum_dtype + ) + t_prefix_logp.unsqueeze(-1) + support_mask = support_mask * teacher_support_mask.to(accum_dtype) + else: + assert teacher_logits is not None + t_final_logits = ( + self._lookup_cp_sharded_teacher_logits( + teacher_logits, + b_idx[:, None], + t_last[:, None], + t_ids, + cp_group, + ) + / temperature + ) + t_final_scores = t_final_logits.to(accum_dtype) - teacher_log_z[ + b_idx, t_last + ].unsqueeze(-1) + t_full_scores = t_final_scores + t_prefix_logp.unsqueeze(-1) + t_full_scores = torch.where( + support_mask > 0, + t_full_scores, + torch.full_like(t_full_scores, NEG_INF), + ) + + # Teacher sparse support pruning is computed without gradients, but the + # student scores must be masked with autograd enabled. + s_full_scores = torch.where( + support_mask > 0, + s_full_scores, + torch.full_like(s_full_scores, NEG_INF), + ) + + # ---- Append REST bucket (per-row valid-mass complement) ---- + s_logp_with_rest = self._append_rest_bucket_logp(s_full_scores) + t_logp_with_rest = self._append_rest_bucket_logp(t_full_scores) + + # Validity: row contributes if it has at least 1 support position. + valid_row = support_mask.sum(dim=-1) >= 1.0 # (K,) bool + + # ---- Batched partition loss ---- + if loss_fn == "bce": + kl_per_pos = self._binary_power_bce_from_logp( + s_logp_with_rest, + t_logp_with_rest, + tau=alm_bce_tau, + ) + elif loss_fn == "jsd": + kl_per_pos = _generalized_jsd( + s_logp_with_rest, + t_logp_with_rest, + jsd_beta, + ) + elif not reverse_kl: + kl_per_pos = torch.nn.functional.kl_div( + s_logp_with_rest, + t_logp_with_rest, + reduction="none", + log_target=True, + ) + else: + kl_per_pos = torch.nn.functional.kl_div( + t_logp_with_rest, + s_logp_with_rest, + reduction="none", + log_target=True, + ) + # Mask padded support columns to zero (REST column stays). + rest_one = torch.ones((K, 1), device=device, dtype=accum_dtype) + full_mask = torch.cat([support_mask, rest_one], dim=-1) + per_row_kl = (kl_per_pos * full_mask).sum(dim=-1) # (K,) + per_row_kl = per_row_kl * valid_row.to(per_row_kl.dtype) + loss_sum = per_row_kl.sum() + + with torch.no_grad(): + s_arg = s_full_scores.argmax(dim=-1) + t_arg = t_full_scores.argmax(dim=-1) + match_row = (s_arg == t_arg) & valid_row + + stats_cpu = torch.stack([valid_row, match_row], dim=0).detach().cpu() + valid_flags = [bool(x) for x in stats_cpu[0].tolist()] + matches = int(stats_cpu[1].sum().item()) + used_count = int(sum(valid_flags)) + effective_support_cpu = support_mask.sum(dim=-1).detach().cpu().tolist() + support_sizes_used = [ + int(size) + for size, is_valid in zip(effective_support_cpu, valid_flags) + if is_valid + ] + + return loss_sum, used_count, matches, support_sizes_used, valid_flags + + def _compute_csd_loss( + self, + i: int, + student_logits: torch.Tensor, + teacher_full_logits: Optional[torch.Tensor], + align: LocalizedAlignment, + *, + teacher_sparse_payload: Optional[SparseTeacherLogits] = None, + teacher_vocab_size: int, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + global_valid_chunks: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, dict[str, Any]]: + """CSD KD term for cross-tokenizer teacher ``i``. + + Prefix-bidir partition KL with MtoM-as-ALM + rest bucket; mismatch + chunks are weighted by ``csd_mismatch_loss_beta``. + ``mtom_as_alm``/``rest_bucket`` are baked in for CSD. + + Per-teacher contract (mirrors ``_compute_same_vocab_kl``): + ``student_logits`` is this rank's CP-contiguous window, + ``teacher_full_logits`` is the reassembled full-vocab teacher window, + and ``align`` carries the localized student/teacher input ids, per-chunk + global position spans, ``pair_valid`` and ``num_chunks``. The body is + native TP/CP-correct: student and dense-teacher rank-3 logits remain + sharded. Exact full-vocabulary log-normalizers are computed in bounded + TP chunks and gathered across CP only after reducing to ``[B,T]``; + arbitrary support values use bounded owner reductions. Production + supplies ``global_valid_chunks`` from the full batch; direct callers + retain a WORLD-reduced fallback. Collectives are no-ops at world size 1 + (single-GPU byte-exact). + + Returns ``(loss, metrics_dict)``. + """ + cfg = self.cfg + device = student_logits.device + # ----- Native TP/CP sharded access ------------------------------- + # Both rank-3 tensors stay sequence-local. Student logits additionally + # stay vocabulary-local; all arbitrary student support accesses below + # use differentiable TP/CP owner reductions. Spans / input_ids arrive + # global from the loss-input adapter. + if teacher_sparse_payload is None: + if teacher_full_logits is None: + raise ValueError("Dense CSD transport is missing teacher logits.") + elif teacher_full_logits is not None: + raise ValueError( + "CSD received both dense and sparse teacher-logit payloads." + ) + # B is the per-sample loop bound; S is the global student-side + # predictor-position bound used in the classify-loop range check. + B, student_local_seq_len, _ = student_logits.shape + student_tp_size = ( + torch.distributed.get_world_size(tp_group) + if tp_group is not None + and torch.distributed.is_initialized() + and torch.distributed.get_world_size(tp_group) > 1 + else 1 + ) + student_cp_size = ( + torch.distributed.get_world_size(cp_group) + if cp_group is not None + and torch.distributed.is_initialized() + and torch.distributed.get_world_size(cp_group) > 1 + else 1 + ) + S = student_local_seq_len * student_cp_size + + # Per-teacher CSD settings. + kl_chunk_shift = bool(cfg.kl_chunk_shift) + temperature = float(cfg.temperature) + reverse_kl = bool(cfg.reverse_kl) + csd_loss_fn = str(cfg.csd_loss_fn or "kl") + csd_jsd_beta = float(cfg.csd_jsd_beta) + csd_alm_bce_tau = float(cfg.csd_alm_bce_tau) + csd_last_pos_loss_fn_cfg = cfg.csd_last_pos_loss_fn + pure_alm = bool(cfg.csd_pure_alm) + realized_only_mismatch = bool(cfg.csd_realized_only_mismatch) + # CSD hybrid: None means "inherit csd_loss_fn"; "kl"/"jsd" overrides + # only the position-N-1 (last-position) partition KL. + effective_last_pos_loss_fn = ( + csd_last_pos_loss_fn_cfg + if csd_last_pos_loss_fn_cfg is not None + else csd_loss_fn + ) + mismatch_loss_beta = float(cfg.csd_mismatch_loss_beta) + noise_filter_topk = int(cfg.csd_noise_filter_topk or 0) + if noise_filter_topk < 0: + raise ValueError( + f"csd_noise_filter_topk must be >= 0, got {noise_filter_topk}" + ) + teacher_topk_ipc_k = int(cfg.teacher_topk_ipc_k) + teacher_topk_keep_realized = teacher_topk_ipc_k > 0 and ( + pure_alm + or realized_only_mismatch + or teacher_topk_ipc_k == 1 + or bool(cfg.teacher_topk_ipc_keep_realized) + ) + # Multi-token chunks use realized-chain probabilities and a rest bucket. + mtom_as_alm = True + + v_t = int(teacher_vocab_size) + teacher_gt_in_topk: Optional[torch.Tensor] = None + teacher_logits: Optional[torch.Tensor] = None + if teacher_sparse_payload is not None: + ( + teacher_topk_logits, + _, + teacher_log_z, + teacher_gt_in_topk, + ) = teacher_sparse_payload + B2, T, _ = teacher_topk_logits.shape + else: + assert teacher_full_logits is not None + # Slice teacher logits to the real tokenizer vocab (drop any padded + # lm_head columns beyond len(tokenizer)) BEFORE indexing. + if teacher_full_logits.shape[-1] > v_t: + teacher_full_logits = teacher_full_logits[..., :v_t] + teacher_logits = teacher_full_logits + # V_t is implicit (common_teacher_idx_t bounds it). Dense teacher + # logits are a contiguous CP-local sequence window, while spans and + # input ids use global positions. + B2, teacher_local_seq_len, _ = teacher_logits.shape + teacher_cp_size = ( + torch.distributed.get_world_size(cp_group) + if cp_group is not None + and torch.distributed.is_initialized() + and torch.distributed.get_world_size(cp_group) > 1 + else 1 + ) + T = teacher_local_seq_len * teacher_cp_size + assert B == B2, f"batch size mismatch: student={B} teacher={B2}" + + # ----- Lazy caches (per-device) ---------------------------------- + if pure_alm: + common_student_idx_t = torch.empty( + (0,), + device=device, + dtype=torch.long, + ) + common_teacher_idx_t = torch.empty( + (0,), + device=device, + dtype=torch.long, + ) + has_common_support = True + else: + common_student_idx_t, common_teacher_idx_t = self._get_common_indices_csd( + device, + i, + ) + has_common_support = common_student_idx_t.numel() > 0 + if not has_common_support: + # Every 1-to-1 chunk would silently fall through to the mismatch + # path, zeroing KL/common and making + # top1_acc_per_chunk meaningless (it is normalized by a common + # count of 0). That looks like a converging run, so fail loudly + # instead: an empty common vocab means the projection matrix or + # subtoks table does not match this (student, teacher) pair. + source = ( + f"student_to_teacher_table_paths[{i}]={self.student_to_teacher_table_paths[i]!r}" + if self.common_indices_from_subtoks + else f"projection_matrix_paths[{i}]={self.projection_matrix_paths[i]!r}" + ) + raise ValueError( + f"Teacher {i}: the exact 1-to-1 common-vocab map is empty, so " + f"no chunk can take the common path. Derived from {source} " + f"(common_indices_from_subtoks=" + f"{self.common_indices_from_subtoks}). Check that this table " + f"was built for this student/teacher tokenizer pair." + ) + # ----- Single-shot host pull of input_ids ------------------------ + # One .detach().cpu().numpy() per tensor, then per-chunk numpy indexing + # inside the classify/collect loop. The ids ride on the localized + # alignment as this CP rank's contiguous window (both unshifted; the + # classify loop applies its own per-chunk kl_chunk_shift); they are + # CP-gathered to the full sequence to match the gathered logits and the + # global spans (no-op at CP=1). + assert align.student_input_ids is not None + assert align.teacher_input_ids is not None + input_ids_student = to_local_if_dtensor(align.student_input_ids) + input_ids_teacher = to_local_if_dtensor(align.teacher_input_ids) + with torch.no_grad(): + input_ids_student = allgather_cp_contiguous_tensor( + input_ids_student, cp_group + ) + input_ids_teacher = allgather_cp_contiguous_tensor( + input_ids_teacher, cp_group + ) + inp_s_np = input_ids_student.detach().cpu().numpy() + inp_t_np = input_ids_teacher.detach().cpu().numpy() + + teacher_gt_in_topk_np: Optional[np.ndarray] = None + teacher_filter_topk_indices_np: Optional[np.ndarray] = None + dense_noise_filter_cp_sharded = False + filter_k = 0 + if noise_filter_topk > 0: + if teacher_sparse_payload is not None: + if teacher_gt_in_topk is None: + raise ValueError( + "csd_noise_filter_topk requires sparse " + "teacher IPC handles with gt_in_topk metadata. Ensure " + "the teacher IPC path receives force_include_token_ids." + ) + teacher_gt_in_topk_np = ( + teacher_gt_in_topk.detach().cpu().numpy().astype(bool) + ) + else: + assert teacher_logits is not None + filter_k = min(noise_filter_topk, int(teacher_logits.shape[-1])) + dense_noise_filter_cp_sharded = teacher_cp_size > 1 + if filter_k > 0 and not dense_noise_filter_cp_sharded: + with torch.no_grad(): + teacher_filter_topk_indices_np = ( + torch.topk( + teacher_logits, + k=filter_k, + dim=-1, + ) + .indices.detach() + .cpu() + .numpy() + ) + + def _teacher_targets_pass_noise_filter( + b_idx: int, + pred_positions: list[int], + label_positions: list[int], + ) -> bool: + if noise_filter_topk <= 0: + return True + if len(pred_positions) != len(label_positions): + return False + if teacher_gt_in_topk_np is not None: + for pred_pos in pred_positions: + if ( + pred_pos < 0 + or pred_pos >= teacher_gt_in_topk_np.shape[1] + or not bool(teacher_gt_in_topk_np[b_idx, pred_pos]) + ): + return False + return True + for pred_pos, label_pos in zip(pred_positions, label_positions): + if pred_pos < 0 or pred_pos >= T: + return False + if label_pos < 0 or label_pos >= inp_t_np.shape[1]: + return False + label_id = int(inp_t_np[b_idx, label_pos]) + if dense_noise_filter_cp_sharded: + assert teacher_logits is not None + found = self._cp_sharded_teacher_target_in_topk( + teacher_logits, + batch_index=b_idx, + global_position=pred_pos, + token_id=label_id, + k=filter_k, + cp_group=cp_group, + ) + elif teacher_filter_topk_indices_np is not None: + found = bool( + np.any( + teacher_filter_topk_indices_np[b_idx, pred_pos] == label_id + ) + ) + else: + found = False + if not found: + return False + return True + + # ----- Alignment payload (padded tensors -> per-chunk records) --- + # The collator ships padded per-chunk position spans + a validity mask + # (dense over max_pairs); a sentinel / OOB pair keeps its slot and is + # zeroed by ``pair_valid``. Spans / pair_valid / num_chunks ride on the + # localized alignment. + assert align.student_spans is not None + assert align.teacher_spans is not None + assert align.pair_valid is not None + assert align.num_chunks is not None + s_spans = to_local_if_dtensor(align.student_spans) # [B, max_pairs, 2] + t_spans = to_local_if_dtensor(align.teacher_spans) # [B, max_pairs, 2] + sample_mask = to_local_if_dtensor(align.sample_mask).bool() # [B] + if sample_mask.ndim != 1 or sample_mask.shape[0] != B: + raise ValueError( + "CSD sample_mask must have shape [B]: " + f"expected ({B},), got {tuple(sample_mask.shape)}." + ) + pair_valid = to_local_if_dtensor(align.pair_valid).bool() # [B, max_pairs] + if pair_valid.ndim != 2 or pair_valid.shape[0] != B: + raise ValueError( + "CSD pair_valid must have shape [B, max_pairs]: " + f"expected first dimension {B}, got {tuple(pair_valid.shape)}." + ) + # Gate the authoritative chunk-validity tensor before classification. + # Every downstream numerator, count, metric, and student gradient is + # derived from the resulting common/mismatch cohorts. + pair_valid = pair_valid & sample_mask.unsqueeze(-1) + num_chunks = to_local_if_dtensor(align.num_chunks) # [B] + # Host-side copies for the classification loop (no GPU sync per + # chunk). + s_spans_np = s_spans.detach().cpu().numpy() + t_spans_np = t_spans.detach().cpu().numpy() + pair_valid_np = pair_valid.detach().cpu().numpy().astype(bool) + num_chunks_np = num_chunks.detach().cpu().numpy().astype(np.int64) + + # ----- Classify chunks ---------------- + common_chunks: list[tuple[int, int, int, int, int]] = [] + mismatch_chunks: list[ + tuple[int, int, int, list[int], list[int], list[int], list[int]] + ] = [] + noise_filtered_common_chunks = 0 + noise_filtered_mismatch_chunks = 0 + T_s_label = input_ids_student.shape[1] + T_t_label = input_ids_teacher.shape[1] + for b in range(B): + n = int(num_chunks_np[b]) + for k in range(n): + if not bool(pair_valid_np[b, k]): + continue + s_start = int(s_spans_np[b, k, 0]) + s_end = int(s_spans_np[b, k, 1]) + t_start = int(t_spans_np[b, k, 0]) + t_end = int(t_spans_np[b, k, 1]) + # Explicit -1 spans are invalid even when pair_valid is set. + if s_start == -1 or t_start == -1: + continue + M = s_end - s_start + N = t_end - t_start + if M <= 0 or N <= 0: + continue + # NOTE: preserve per-chunk kl_chunk_shift. + # Do NOT reuse nemo_rl's tensor-level global shift — + # CSD needs the conditional per-chunk shift (chunks at + # sequence-start are NOT shifted). + if kl_chunk_shift and s_start > 0 and t_start > 0: + s_pred = list(range(s_start - 1, s_end - 1)) + t_pred = list(range(t_start - 1, t_end - 1)) + s_labels = list(range(s_start, s_end)) + t_labels = list(range(t_start, t_end)) + else: + s_pred = list(range(s_start, s_end)) + t_pred = list(range(t_start, t_end)) + s_labels = s_pred + t_labels = t_pred + + if any(pos < 0 or pos >= S for pos in s_pred): + continue + if any(pos < 0 or pos >= T for pos in t_pred): + continue + if any(pos < 0 or pos >= T_s_label for pos in s_labels): + continue + if any(pos < 0 or pos >= T_t_label for pos in t_labels): + continue + + if M == 1 and N == 1 and (has_common_support or pure_alm): + if not _teacher_targets_pass_noise_filter( + b, + t_pred, + t_labels, + ): + noise_filtered_common_chunks += 1 + continue + common_chunks.append( + ( + b, + s_pred[0], + t_pred[0], + int(inp_s_np[b, s_labels[0]]), + int(inp_t_np[b, t_labels[0]]), + ) + ) + else: + if not _teacher_targets_pass_noise_filter( + b, + t_pred, + t_labels, + ): + noise_filtered_mismatch_chunks += 1 + continue + mismatch_chunks.append( + (b, M, N, s_pred, t_pred, s_labels, t_labels) + ) + + # ----- No early-out on empty chunks ----------------------------- + # A microbatch with no valid chunks (empty common + mismatch) is NOT + # returned early: it must still reach the ``global_valid_chunks`` + # WORLD-reduce below so a rank with zero local chunks participates in + # the collective (other DP/CP ranks may have chunks). The empty + # common/mismatch loops leave both partition losses at zero and the + # combine step returns a gradient-connected zero. + accum_dtype = ( + torch.float32 + if student_logits.dtype in (torch.float16, torch.bfloat16) + else student_logits.dtype + ) + common_loss = torch.tensor(0.0, device=device, dtype=accum_dtype) + mismatch_loss = torch.tensor(0.0, device=device, dtype=accum_dtype) + common_matches = 0 + mismatch_matches = 0 + + # ----- Pre-compute log_Z's per (b, predictor-position) ---------- + # Partition-KL (rest_bucket=True) needs log Z[b, s] for every + # (b, predictor) position the chunk loops reference. Chunked along + # the seq axis to bound peak fp32 temporary. + def _precompute_local_log_z(_logits: torch.Tensor, _chunk: int = 256): + _B, _S, _V = _logits.shape + _out = torch.empty( + (_B, _S), + dtype=accum_dtype, + device=_logits.device, + ) + for _s in range(0, _S, _chunk): + _end = min(_s + _chunk, _S) + _chunk_t = _logits[:, _s:_end, :].to(accum_dtype) / temperature + _out[:, _s:_end] = torch.logsumexp(_chunk_t, dim=-1) + return _out + + student_log_z_local = chunked_vocab_parallel_logsumexp( + student_logits, + temperature, + tp_group=tp_group, + chunk_size=CSD_STUDENT_LOSS_CHUNK_SIZE, + ) + # Only the scalar normalizer per student row crosses CP. Backward of + # this contiguous gather SUMs the replicated CP gradient, cancelling + # the outer loss / CP scaling exactly as the deleted rank-3 gather did. + student_log_z = allgather_cp_contiguous_tensor( + student_log_z_local, + cp_group, + ) + if teacher_sparse_payload is None: + assert teacher_logits is not None + with torch.no_grad(): + teacher_log_z_local = _precompute_local_log_z(teacher_logits) + # Only the scalar normalizer per teacher row is materialized + # across CP. The multi-GiB dense logits stay sequence-sharded. + teacher_log_z = allgather_cp_contiguous_tensor( + teacher_log_z_local, + cp_group, + ) + + # ----- Common-vocab KL ---------------- + if common_chunks: + sb_full = torch.tensor( + [c[0] for c in common_chunks], + device=device, + dtype=torch.long, + ) + sp_full = torch.tensor( + [c[1] for c in common_chunks], + device=device, + dtype=torch.long, + ) + tp_full = torch.tensor( + [c[2] for c in common_chunks], + device=device, + dtype=torch.long, + ) + rs_full = torch.tensor( + [c[3] for c in common_chunks], + device=device, + dtype=torch.long, + ) + rt_full = torch.tensor( + [c[4] for c in common_chunks], + device=device, + dtype=torch.long, + ) + + common_loss_sum = torch.tensor( + 0.0, + device=device, + dtype=accum_dtype, + ) + common_matches_sum_t = torch.zeros( + (), + device=device, + dtype=torch.long, + ) + + _env_common_mb = os.environ.get("NRL_CSD_COMMON_LOOP_MB") + common_mb = int(_env_common_mb) if _env_common_mb else len(common_chunks) + if student_tp_size > 1 or student_cp_size > 1: + # Student owner lookups issue TP/CP collectives. Fix the row + # bound in code so environment drift cannot change collective + # counts or order across ranks, regardless of teacher transport. + common_mb = min(len(common_chunks), 64) + elif teacher_sparse_payload is not None and not _env_common_mb: + common_mb = min(common_mb, 64) + teacher_to_common_student = None + if teacher_sparse_payload is not None and not pure_alm: + teacher_to_common_student = self._get_csd_teacher_to_common_student( + device, + i, + v_t, + common_student_idx_t, + common_teacher_idx_t, + ) + + for mb_start in range(0, len(common_chunks), common_mb): + mb_end = min(mb_start + common_mb, len(common_chunks)) + sb = sb_full[mb_start:mb_end] + sp = sp_full[mb_start:mb_end] + tp = tp_full[mb_start:mb_end] + rs = rs_full[mb_start:mb_end] + rt = rt_full[mb_start:mb_end] + support_mask = None + common_valid_row = None + if pure_alm: + student_common_logp = ( + self._lookup_tp_cp_sharded_student_logits( + student_logits, + sb, + sp, + rs, + tp_group=tp_group, + cp_group=cp_group, + ).to(accum_dtype) + / temperature + - student_log_z[sb, sp] + ).unsqueeze(-1) + if teacher_sparse_payload is None: + assert teacher_logits is not None + teacher_common_logp = ( + self._lookup_cp_sharded_teacher_logits( + teacher_logits, + sb, + tp, + rt, + cp_group, + ).to(accum_dtype) + / temperature + - teacher_log_z[sb, tp] + ).unsqueeze(-1) + common_valid_row = torch.ones( + (sb.shape[0],), + device=device, + dtype=torch.bool, + ) + else: + teacher_topk_logits, teacher_topk_indices, _, _ = ( + teacher_sparse_payload + ) + t_real_logp, t_real_found = self._lookup_sparse_teacher_logp( + teacher_topk_logits, + teacher_topk_indices, + teacher_log_z, + sb, + tp, + rt, + temperature, + ) + teacher_common_logp = t_real_logp.to(accum_dtype).unsqueeze(-1) + common_valid_row = t_real_found + support_mask = common_valid_row.unsqueeze(-1) + neg_inf = torch.full_like(student_common_logp, -1.0e30) + student_common_logp = torch.where( + support_mask, + student_common_logp, + neg_inf, + ) + teacher_common_logp = torch.where( + support_mask, + teacher_common_logp, + neg_inf, + ) + elif teacher_sparse_payload is None: + assert teacher_logits is not None + student_common_logits = ( + self._lookup_tp_cp_sharded_student_logits( + student_logits, + sb[:, None], + sp[:, None], + common_student_idx_t[None, :], + tp_group=tp_group, + cp_group=cp_group, + ) + / temperature + ) + teacher_common_logits = self._lookup_cp_sharded_teacher_logits( + teacher_logits, + sb[:, None], + tp[:, None], + common_teacher_idx_t[None, :], + cp_group, + ) + teacher_common_logits = teacher_common_logits / temperature + # Append the complementary probability mass as a rest bucket. + s_log_z_mb = student_log_z[sb, sp].unsqueeze(-1) + t_log_z_mb = teacher_log_z[sb, tp].unsqueeze(-1) + student_common_logp = ( + student_common_logits.to(accum_dtype) - s_log_z_mb + ) + teacher_common_logp = ( + teacher_common_logits.to(accum_dtype) - t_log_z_mb + ) + else: + assert teacher_to_common_student is not None + teacher_topk_logits, teacher_topk_indices, _, _ = ( + teacher_sparse_payload + ) + topk_ids = teacher_topk_indices[sb, tp] + topk_vals = teacher_topk_logits[sb, tp] + mapped_student = teacher_to_common_student[topk_ids.long()] + support_mask = mapped_student >= 0 + if teacher_topk_keep_realized: + topk_width = int(topk_ids.shape[-1]) + alt_budget = max(topk_width - 1, 0) + support_mask = ( + support_mask + & (mapped_student != rs[:, None]) + & (topk_ids.long() != rt[:, None]) + ) + support_mask = self._topk_mask_by_score( + support_mask, + topk_vals.to(accum_dtype), + alt_budget, + ) + safe_student_ids = mapped_student.clamp_min(0) + student_common_logits = ( + self._lookup_tp_cp_sharded_student_logits( + student_logits, + sb[:, None], + sp[:, None], + safe_student_ids, + tp_group=tp_group, + cp_group=cp_group, + ) + / temperature + ) + student_common_logp = student_common_logits.to( + accum_dtype + ) - student_log_z[sb, sp].unsqueeze(-1) + teacher_common_logp = topk_vals.to( + accum_dtype + ) / temperature - teacher_log_z[sb, tp].unsqueeze(-1) + neg_inf = torch.full_like(student_common_logp, -1.0e30) + student_common_logp = torch.where( + support_mask, + student_common_logp, + neg_inf, + ) + teacher_common_logp = torch.where( + support_mask, + teacher_common_logp, + neg_inf, + ) + if teacher_topk_keep_realized: + t_real_logp, t_real_found = self._lookup_sparse_teacher_logp( + teacher_topk_logits, + teacher_topk_indices, + teacher_log_z, + sb, + tp, + rt, + temperature, + ) + realized_common_student = teacher_to_common_student[rt.long()] + realized_common_mask = t_real_found & ( + realized_common_student == rs + ) + s_real_logp = ( + self._lookup_tp_cp_sharded_student_logits( + student_logits, + sb, + sp, + rs, + tp_group=tp_group, + cp_group=cp_group, + ).to(accum_dtype) + / temperature + - student_log_z[sb, sp] + ) + student_common_logp = torch.cat( + [ + student_common_logp, + torch.where( + realized_common_mask, + s_real_logp, + torch.full_like(s_real_logp, -1.0e30), + ).unsqueeze(-1), + ], + dim=-1, + ) + teacher_common_logp = torch.cat( + [ + teacher_common_logp, + torch.where( + realized_common_mask, + t_real_logp.to(accum_dtype), + torch.full_like(t_real_logp, -1.0e30), + ).unsqueeze(-1), + ], + dim=-1, + ) + support_mask = torch.cat( + [ + support_mask, + realized_common_mask.unsqueeze(-1), + ], + dim=-1, + ) + common_valid_row = support_mask.any(dim=-1) + student_loss_logp = self._append_rest_bucket_logp( + student_common_logp, + ) + teacher_loss_logp = self._append_rest_bucket_logp( + teacher_common_logp, + ) + if csd_loss_fn == "bce": + per_elem = self._binary_power_bce_from_logp( + student_loss_logp, + teacher_loss_logp, + tau=csd_alm_bce_tau, + ) + elif csd_loss_fn == "jsd": + per_elem = _generalized_jsd( + student_loss_logp, + teacher_loss_logp, + csd_jsd_beta, + ) + elif not reverse_kl: + per_elem = torch.nn.functional.kl_div( + student_loss_logp, + teacher_loss_logp, + reduction="none", + log_target=True, + ) + else: + per_elem = torch.nn.functional.kl_div( + teacher_loss_logp, + student_loss_logp, + reduction="none", + log_target=True, + ) + if support_mask is not None: + rest_one = torch.ones( + (support_mask.shape[0], 1), + device=device, + dtype=per_elem.dtype, + ) + full_mask = torch.cat( + [support_mask.to(per_elem.dtype), rest_one], + dim=-1, + ) + per_row = (per_elem * full_mask).sum(dim=-1) + per_row = per_row * common_valid_row.to(per_row.dtype) + common_loss_sum = common_loss_sum + per_row.sum().to( + accum_dtype, + ) + else: + common_loss_sum = common_loss_sum + per_elem.sum().to( + accum_dtype, + ) + with torch.no_grad(): + row_match = student_common_logp.argmax( + dim=-1 + ) == teacher_common_logp.argmax(dim=-1) + if common_valid_row is not None: + row_match = row_match & common_valid_row + common_matches_sum_t = common_matches_sum_t + row_match.sum().to( + torch.long + ) + common_loss = common_loss_sum / float(len(common_chunks)) + common_matches = int(common_matches_sum_t.item()) + + # ----- Mismatch loss ------------------ + skipped_no_support = 0 + support_sizes: list[int] = [] + prefix_index = None + if mismatch_chunks: + if pure_alm or realized_only_mismatch: + _prefix_index_prededuped = False + max_chain_t = 0 + max_chain_s = 0 + else: + prefix_index = self._ensure_bidir_prefix_support_index(device, i) + _prefix_index_prededuped = bool(prefix_index.get("_prededuped", False)) + max_chain_t = int(prefix_index.get("max_chain_t", 0)) + max_chain_s = int(prefix_index.get("max_chain_s", 0)) + + chunk_records: list[ + tuple[ + int, + int, + int, + list[int], + list[int], + list[int], + list[int], + list[int], + list[int], + int, + int, + ] + ] = [] + for b, M, N, s_pred, t_pred, s_labels, t_labels in mismatch_chunks: + if pure_alm or realized_only_mismatch: + # Realized token chains against REST. + s_realized_id = int(inp_s_np[b, s_labels[-1]]) + t_realized_id = int(inp_t_np[b, t_labels[-1]]) + s_prefix_pred = list(s_pred[:-1]) + s_prefix_lab_ids = [int(inp_s_np[b, pos]) for pos in s_labels[:-1]] + t_prefix_pred = list(t_pred[:-1]) + t_prefix_lab_ids = [int(inp_t_np[b, pos]) for pos in t_labels[:-1]] + chunk_records.append( + ( + b, + int(s_pred[-1]), + int(t_pred[-1]), + s_prefix_pred, + s_prefix_lab_ids, + t_prefix_pred, + t_prefix_lab_ids, + [s_realized_id], + [t_realized_id], + s_realized_id, + t_realized_id, + ) + ) + elif M == 1 and N > 1: + # 1-to-many + s_realized_id = int(inp_s_np[b, s_labels[0]]) + t_realized_id = int(inp_t_np[b, t_labels[-1]]) + if N > max_chain_t and not teacher_topk_keep_realized: + skipped_no_support += 1 + continue + pairs = () + if N <= max_chain_t: + assert prefix_index is not None + teacher_prefix = tuple( + int(inp_t_np[b, pos]) for pos in t_labels[: N - 1] + ) + pairs = prefix_index["forward"].get( + (N, teacher_prefix), + (), + ) + if not pairs and not teacher_topk_keep_realized: + skipped_no_support += 1 + continue + s_list, t_list = self._unique_bidir_pairs_cpu( + pairs, + swap=False, + assume_unique=_prefix_index_prededuped, + ) + if teacher_topk_keep_realized and ( + s_realized_id, + t_realized_id, + ) not in set(zip(s_list, t_list)): + s_list.append(s_realized_id) + t_list.append(t_realized_id) + min_support = 1 if teacher_topk_keep_realized else 2 + if len(s_list) < min_support: + skipped_no_support += 1 + continue + t_prefix_pred = list(t_pred[:-1]) + t_prefix_lab_ids = [int(inp_t_np[b, pos]) for pos in t_labels[:-1]] + chunk_records.append( + ( + b, + int(s_pred[0]), + int(t_pred[-1]), + [], + [], + t_prefix_pred, + t_prefix_lab_ids, + s_list, + t_list, + s_realized_id, + t_realized_id, + ) + ) + elif M > 1 and N == 1: + # Many-to-1 + s_realized_id = int(inp_s_np[b, s_labels[-1]]) + t_realized_id = int(inp_t_np[b, t_labels[0]]) + if M > max_chain_s and not teacher_topk_keep_realized: + skipped_no_support += 1 + continue + pairs = () + if M <= max_chain_s: + assert prefix_index is not None + student_prefix = tuple( + int(inp_s_np[b, pos]) for pos in s_labels[: M - 1] + ) + pairs = prefix_index["reverse"].get( + (M, student_prefix), + (), + ) + if not pairs and not teacher_topk_keep_realized: + skipped_no_support += 1 + continue + s_list, t_list = self._unique_bidir_pairs_cpu( + pairs, + swap=True, + assume_unique=_prefix_index_prededuped, + ) + if teacher_topk_keep_realized and ( + s_realized_id, + t_realized_id, + ) not in set(zip(s_list, t_list)): + s_list.append(s_realized_id) + t_list.append(t_realized_id) + min_support = 1 if teacher_topk_keep_realized else 2 + if len(s_list) < min_support: + skipped_no_support += 1 + continue + s_prefix_pred = list(s_pred[:-1]) + s_prefix_lab_ids = [int(inp_s_np[b, pos]) for pos in s_labels[:-1]] + chunk_records.append( + ( + b, + int(s_pred[-1]), + int(t_pred[0]), + s_prefix_pred, + s_prefix_lab_ids, + [], + [], + s_list, + t_list, + s_realized_id, + t_realized_id, + ) + ) + else: + # Many-to-many chunks use the two realized token chains. + assert mtom_as_alm, "CSD requires realized-chain support" + s_realized_id = int(inp_s_np[b, s_labels[-1]]) + t_realized_id = int(inp_t_np[b, t_labels[-1]]) + s_list = [s_realized_id] + t_list = [t_realized_id] + s_prefix_pred = list(s_pred[:-1]) + s_prefix_lab_ids = [int(inp_s_np[b, pos]) for pos in s_labels[:-1]] + t_prefix_pred = list(t_pred[:-1]) + t_prefix_lab_ids = [int(inp_t_np[b, pos]) for pos in t_labels[:-1]] + chunk_records.append( + ( + b, + int(s_pred[-1]), + int(t_pred[-1]), + s_prefix_pred, + s_prefix_lab_ids, + t_prefix_pred, + t_prefix_lab_ids, + s_list, + t_list, + s_realized_id, + t_realized_id, + ) + ) + + if chunk_records: + ( + loss_sum_b, + used_b, + matches_b, + support_sizes_b, + _valid_flags_b, + ) = self._partition_kl_mismatch_batched( + chunk_records, + student_logits, + teacher_logits, + student_log_z, + teacher_log_z, + temperature, + reverse_kl, + loss_fn=effective_last_pos_loss_fn, + jsd_beta=csd_jsd_beta, + alm_bce_tau=csd_alm_bce_tau, + teacher_sparse=( + None + if teacher_sparse_payload is None + else ( + teacher_sparse_payload[0], + teacher_sparse_payload[1], + ) + ), + teacher_sparse_keep_realized=teacher_topk_keep_realized, + tp_group=tp_group, + cp_group=cp_group, + ) + if used_b > 0: + mismatch_loss = loss_sum_b / float(used_b) + mismatch_matches += matches_b + support_sizes = support_sizes_b + + # ----- Combine + T^2 scale ----------- + effective_mismatch_count = len(support_sizes) + effective_count = len(common_chunks) + effective_mismatch_count + # DP×CP-global valid-chunk count, computed in-loss. Reduce over the full + # mesh (WORLD), then divide out the TP replication: the chunk count is + # identical on every TP rank (chunks are vocab-independent), so + # ``WORLD-sum == tp_world × (DP×CP-sum)`` exactly, and dividing by + # ``tp_world`` recovers the DP×CP count. This matches the CE term's + # ``global_valid_toks`` (which carries no TP factor), so the KD keeps its + # configured weight relative to CE at TP>1 — unlike the deleted P-KL/gold + # WORLD-reduce, which under-weighted the KD by 1/tp. The reduce is + # unconditional so a no-chunk rank still fires the collective (the + # early-return was removed for exactly this). At world size 1 (and CPU + # parity) it is a no-op, so ``global_valid_chunks`` collapses to the local + # ``effective_count`` and the normalization is byte-exact. + if global_valid_chunks is None: + local_count = torch.tensor( + float(effective_count), device=device, dtype=torch.float32 + ) + if self._dp_cp_group is not None: + # DP x CP group: it already excludes the TP replicas (so there is + # no TP factor to divide out) and, crucially, it lives inside a + # single pipeline stage. A WORLD reduce here deadlocks under PP, + # because only the last stage runs the loss and joins it. + global_valid_chunks = group_all_reduce_sum( + local_count, group=self._dp_cp_group + ) + else: + global_valid_chunks = group_all_reduce_sum( + local_count, group=torch.distributed.group.WORLD + ) + tp_world = ( + torch.distributed.get_world_size(tp_group) + if tp_group is not None + else 1 + ) + if tp_world > 1: + global_valid_chunks = global_valid_chunks / float(tp_world) + mismatch_last_loss = mismatch_loss + mismatch_loss = mismatch_last_loss * mismatch_loss_beta + # The mismatch objective retains its full coefficient after aggregation. + if effective_count > 0: + # Sum of all chunk-level KLs for this microbatch. common_loss / + # mismatch_loss are local per-chunk averages over their + # respective partitions; multiplying by chunk counts recovers + # the un-normalised partition sums. + chunk_sum = common_loss * float(len(common_chunks)) + mismatch_loss * float( + effective_mismatch_count + ) + # Per-chunk normalisation. ``global_valid_chunks`` is either the + # in-loss WORLD-reduce above or a caller-supplied count (both + # >= this rank's ``effective_count`` > 0 here, so no div-by-zero). + # The trainer's sum-reduction across microbatches and ranks then + # yields `total_chunk_KL × T² / global_valid_chunks` — the natural + # per-chunk KL averaged over the global batch, matching the + # per-token convention used by _compute_ce (each loss normalises by + # its own natural global denominator). + loss_total = chunk_sum / global_valid_chunks.to(chunk_sum.dtype) + else: + # No usable chunks. Keep gradient-connected zero. + loss_total = (student_logits.sum() * 0.0).to(accum_dtype) + top1_accuracy = (common_matches + mismatch_matches) / max( + effective_count, + 1, + ) + + # Final loss × T^2. + final_loss = loss_total * (temperature**2) + + # Metric naming convention for wandb clarity: + # - `loss` / `train:loss`: sum-reduced by the trainer. When + # global_valid_chunks is provided, final wandb value is the + # per-chunk KL averaged over the global batch + # (`total_chunk_KL × T² / global_valid_chunks`). The dispatcher + # additionally overwrites `loss` with the dynamic-scaled + # `kl_scale * kl + ce` combination — that final value mixes + # a per-chunk KL term and a per-token CE term (consistent + # gradient balance by construction of kl_scale). `train:loss` + # is also the checkpoint-selection metric. + # - `*_per_chunk`: mean-reduced by the trainer (listed in + # `xtoken_off_policy_distillation.MEAN_REDUCED_MB_METRICS`); final + # wandb value is the average per-mb local chunk-average — a clean + # per-chunk quantity when chunk counts are roughly uniform across + # microbatches. + # - `num_*`: sum-reduced; total count across the step. + metrics: dict[str, Any] = { + "loss": final_loss.item(), + "train:loss": final_loss.item(), + "kl_common_per_chunk": common_loss.item(), + "kl_partition_last_per_chunk": mismatch_last_loss.item(), + "kl_mismatch_scaled_per_chunk": mismatch_loss.item(), + "num_common_chunks": len(common_chunks), + "num_mismatch_chunks": effective_mismatch_count, + "num_noise_filtered_common_chunks": noise_filtered_common_chunks, + "num_noise_filtered_mismatch_chunks": noise_filtered_mismatch_chunks, + "top1_acc_per_chunk": float(top1_accuracy), + "num_valid_samples": sample_mask.sum().item(), + } + metrics["csd_mismatch_loss_beta"] = mismatch_loss_beta + if noise_filter_topk > 0: + metrics["csd_noise_filter_topk"] = noise_filter_topk + return final_loss, metrics + + # ------------------------------------------------------------------ # + # Owner-local evaluation (Megatron students) + # ------------------------------------------------------------------ # + @staticmethod + def _detached_cp_sum( + value: torch.Tensor, cp_group: Optional[torch.distributed.ProcessGroup] + ) -> torch.Tensor: + """Detached SUM of a CP-partial scalar, identical on every CP rank.""" + reported = value.detach().clone() + if ( + cp_group is not None + and torch.distributed.is_initialized() + and torch.distributed.get_world_size(cp_group) > 1 + ): + torch.distributed.all_reduce( + reported, op=torch.distributed.ReduceOp.SUM, group=cp_group + ) + return reported + + @staticmethod + def _raise_if_any_rank_failed( + error: Optional[Exception], + *, + device: torch.device, + tp_group: Optional[torch.distributed.ProcessGroup], + cp_group: Optional[torch.distributed.ProcessGroup], + ) -> None: + """Make every TP x CP rank raise together when one rank's setup failed. + + Reducing an ok flag with MIN over the TP group and then the CP group + reaches every rank of the TP x CP grid, so a failure on one rank cannot + leave the others waiting in the next collective. + """ + ok = torch.tensor([0 if error is not None else 1], device=device) + for group in (tp_group, cp_group): + if ( + group is not None + and torch.distributed.is_initialized() + and torch.distributed.get_world_size(group) > 1 + ): + torch.distributed.all_reduce( + ok, op=torch.distributed.ReduceOp.MIN, group=group + ) + if error is not None: + raise error + if int(ok.item()) == 0: + raise RuntimeError( + "Owner-local CSD evaluation failed on another tensor- or " + "context-parallel rank; see that rank's error." + ) + + def _compute_owner_local_ce( + self, + logits: torch.Tensor, + data: BatchedDataDict[CSDLossDataDict], + global_valid_toks: torch.Tensor, + *, + tp_group: Optional[torch.distributed.ProcessGroup], + cp_group: Optional[torch.distributed.ProcessGroup], + ) -> torch.Tensor: + """CP-partial next-token CE over this rank's own student rows. + + Each CP rank sums ``-log p(label)`` over the predictor rows it holds and + divides by the global token count, so the CP sum of the partials is the + full CE. Only selected log-probabilities cross TP; nothing is gathered + across CP. + """ + full_ids = to_local_if_dtensor(data["input_ids"]) + full_mask = to_local_if_dtensor(data["token_mask"]) + sample_mask = to_local_if_dtensor(data["sample_mask"]) + positions = cp_load_balanced_global_positions( + int(logits.shape[1]), cp_group, device=logits.device + ) + predictor = positions < int(full_ids.shape[1]) - 1 + local_rows = torch.arange( + positions.numel(), device=logits.device, dtype=torch.long + )[predictor] + batch_size = int(full_ids.shape[0]) + batch_indices = torch.arange( + batch_size, device=logits.device + ).repeat_interleave(local_rows.numel()) + seq_indices = local_rows.repeat(batch_size) + label_positions = (positions[predictor] + 1).repeat(batch_size) + tp_rank = ( + torch.distributed.get_rank(tp_group) + if tp_group is not None and torch.distributed.is_initialized() + else 0 + ) + local_vocab = int(logits.shape[-1]) + logp = distributed_selected_logprobs( + logits, + batch_indices, + seq_indices, + full_ids.to(logits.device)[batch_indices, label_positions], + vocab_start_index=tp_rank * local_vocab, + vocab_end_index=(tp_rank + 1) * local_vocab, + real_vocab_size=int(self.student_vocab_size), + temperature=1.0, + tp_group=tp_group, + row_chunk_size=self.cfg.csd_selected_row_chunk_size, + ) + mask = full_mask.to(logits.device)[batch_indices, label_positions].to( + logp.dtype + ) * sample_mask.to(logits.device)[batch_indices].to(logp.dtype) + return -(logp * mask).sum() / global_valid_toks.to(logp.dtype) + + @staticmethod + def _owner_local_next_token_accuracy( + logits: torch.Tensor, + data: BatchedDataDict[CSDLossDataDict], + *, + tp_group: Optional[torch.distributed.ProcessGroup], + cp_group: Optional[torch.distributed.ProcessGroup], + ) -> torch.Tensor: + """Next-token top-1 accuracy over CP-local rows (scalar collectives only).""" + with torch.no_grad(): + full_ids = to_local_if_dtensor(data["input_ids"]).to(logits.device) + full_mask = to_local_if_dtensor(data["token_mask"]).to(logits.device) + sample_mask = to_local_if_dtensor(data["sample_mask"]).to(logits.device) + positions = cp_load_balanced_global_positions( + int(logits.shape[1]), cp_group, device=logits.device + ) + predictor = positions < int(full_ids.shape[1]) - 1 + local_rows = torch.arange( + positions.numel(), device=logits.device, dtype=torch.long + )[predictor] + label_positions = positions[predictor] + 1 + prediction = vocab_parallel_argmax(logits, tp_group=tp_group) + mask = ( + full_mask[:, label_positions].float() + * sample_mask.unsqueeze(-1).float() + ) + correct = ( + (prediction[:, local_rows] == full_ids[:, label_positions]).float() + * mask + ).sum() + stats = torch.stack((correct, mask.sum())) + if ( + cp_group is not None + and torch.distributed.is_initialized() + and torch.distributed.get_world_size(cp_group) > 1 + ): + torch.distributed.all_reduce(stats, group=cp_group) + return stats[0] / stats[1].clamp(min=1.0) + + def _owner_local_sum_kd( + self, + logits: torch.Tensor, + student_logits_contig: Optional[torch.Tensor], + teacher_full_logits_by_idx: dict[int, torch.Tensor], + teacher_row_payloads_by_idx: dict[int, TeacherRowPayload], + aligns_by_idx: dict[int, LocalizedAlignment], + kd_normalizer: torch.Tensor, + global_valid_chunks_by_idx: Optional[dict[int, torch.Tensor]], + *, + tp_group: Optional[torch.distributed.ProcessGroup], + cp_group: Optional[torch.distributed.ProcessGroup], + ) -> tuple[torch.Tensor, dict[str, Any]]: + """Static-weight sum of CP-partial per-teacher KD terms. + + Cross-tokenizer teachers run :meth:`_compute_csd_loss_owner_local`. A + same-tokenizer teacher keeps the contiguous-window KL, which is already + a per-window partial over the global token count. Reported per-teacher + values are CP sums. + """ + device = logits.device + if self.normalize_teacher_by_vocab: + min_log_vocab = torch.log( + torch.tensor(float(min(self.teacher_vocab_sizes)), device=device) + ) + total_kd: Optional[torch.Tensor] = None + per_metrics: dict[str, Any] = {} + # Deterministic teacher order: each KD term issues its own collectives. + for i in range(self.num_teachers): + if self._teacher_is_same_vocab(i): + if student_logits_contig is None or i not in teacher_full_logits_by_idx: + raise ValueError( + f"Same-vocab teacher {i} requires full teacher logits and " + "the contiguous student window." + ) + kd_i = same_tokenizer_topk_kl( + student_logits_contig, + teacher_full_logits_by_idx[i], + kd_mask=self._student_kd_mask(aligns_by_idx[i]), + sample_mask=aligns_by_idx[i].sample_mask, + num_kd_tokens=kd_normalizer, + vocab_size=self.student_vocab_size, + vocab_topk=self.vocab_topk, + temperature=self.temperature, + reverse_kl=self.reverse_kl, + tp_group=tp_group, + cp_group=cp_group, + ) + m_i: dict[str, Any] = { + "kl_loss": self._detached_cp_sum(kd_i, cp_group).item() + } + else: + if ( + global_valid_chunks_by_idx is None + or i not in global_valid_chunks_by_idx + ): + raise ValueError( + "Owner-local CSD needs the step-global valid-chunk count " + f"for teacher {i}." + ) + if i not in teacher_row_payloads_by_idx: + raise ValueError( + f"Cross-tokenizer teacher {i} has no IPC row payload." + ) + kd_i, m_i = self._compute_csd_loss_owner_local( + i, + logits, + teacher_row_payloads_by_idx[i], + aligns_by_idx[i], + teacher_vocab_size=self.teacher_vocab_sizes[i], + global_valid_chunks=global_valid_chunks_by_idx[i], + tp_group=tp_group, + cp_group=cp_group, + ) + weight = self.teacher_weights[i] + if self.normalize_teacher_by_vocab: + weight = weight * float( + torch.log( + torch.tensor(float(self.teacher_vocab_sizes[i]), device=device) + ) + / min_log_vocab + ) + weighted = kd_i * weight + total_kd = weighted if total_kd is None else total_kd + weighted + per_metrics[f"teacher_{i}/weighted_kl"] = float(m_i["kl_loss"] * weight) + for key in ( + "kl_loss", + "kl_common_per_chunk", + "kl_partition_last_per_chunk", + "top1_acc_per_chunk", + "num_common_chunks", + "num_mismatch_chunks", + ): + if key in m_i: + per_metrics[f"{key}_t{i}"] = m_i[key] + per_metrics[f"weight_t{i}"] = float(self.teacher_weights[i]) + assert total_kd is not None + return total_kd, per_metrics + + def _partition_divergence( + self, + student_logp: torch.Tensor, + teacher_logp: torch.Tensor, + loss_fn_name: str, + ) -> torch.Tensor: + """Per-element partition divergence for one CSD loss setting.""" + if loss_fn_name == "bce": + return self._binary_power_bce_from_logp( + student_logp, teacher_logp, tau=float(self.cfg.csd_alm_bce_tau) + ) + if loss_fn_name == "jsd": + return _generalized_jsd( + student_logp, teacher_logp, float(self.cfg.csd_jsd_beta) + ) + if not self.reverse_kl: + return torch.nn.functional.kl_div( + student_logp, teacher_logp, reduction="none", log_target=True + ) + return torch.nn.functional.kl_div( + teacher_logp, student_logp, reduction="none", log_target=True + ) + + def _owner_local_mismatch_support( + self, + record: tuple[int, int, int, list[int], list[int], list[int], list[int], bool], + *, + inp_s_np: np.ndarray, + inp_t_np: np.ndarray, + prefix_index: Optional[dict[str, Any]], + keep_realized: bool, + ) -> Optional[tuple[list[int], list[int], int, int]]: + """``(s_ids, t_ids, realized_s, realized_t)`` for one mismatch chunk. + + Mirrors the replicated CSD rules: realized-only chunks and many-to-many + chunks use the realized pair; 1-to-many and many-to-1 chunks use the + prefix-support alternatives, plus the realized pair when the sparse + teacher keeps realized labels. Returns None for an unsupported chunk. + """ + b, M, N, _s_pred, s_labels, _t_pred, t_labels, _owned = record + realized_s = int(inp_s_np[b, s_labels[-1]]) + realized_t = int(inp_t_np[b, t_labels[-1]]) + if ( + self.cfg.csd_realized_only_mismatch + or prefix_index is None + or (M > 1 and N > 1) + ): + return [realized_s], [realized_t], realized_s, realized_t + if M == 1: + max_chain = int(prefix_index.get("max_chain_t", 0)) + length, side, swap = N, "forward", False + prefix = tuple(int(inp_t_np[b, pos]) for pos in t_labels[: N - 1]) + else: + max_chain = int(prefix_index.get("max_chain_s", 0)) + length, side, swap = M, "reverse", True + prefix = tuple(int(inp_s_np[b, pos]) for pos in s_labels[: M - 1]) + if length > max_chain and not keep_realized: + return None + pairs = ( + prefix_index[side].get((length, prefix), ()) if length <= max_chain else () + ) + if not pairs and not keep_realized: + return None + s_list, t_list = self._unique_bidir_pairs_cpu( + pairs, + swap=swap, + assume_unique=bool(prefix_index.get("_prededuped", False)), + ) + if keep_realized and (realized_s, realized_t) not in set(zip(s_list, t_list)): + s_list.append(realized_s) + t_list.append(realized_t) + if len(s_list) < (1 if keep_realized else 2): + return None + return s_list, t_list, realized_s, realized_t + + def _compute_csd_loss_owner_local( + self, + i: int, + student_logits: torch.Tensor, + payload: TeacherRowPayload, + align: LocalizedAlignment, + *, + teacher_vocab_size: int, + global_valid_chunks: torch.Tensor, + tp_group: Optional[torch.distributed.ProcessGroup], + cp_group: Optional[torch.distributed.ProcessGroup], + ) -> tuple[torch.Tensor, dict[str, Any]]: + """CP-partial CSD KD term for cross-tokenizer teacher ``i``. + + The student logits stay ``[B, S_pad/CP, V/TP]`` in their load-balanced + CP layout. Every rank classifies all chunks the same way, but only the + rank holding a chunk's final student predictor row evaluates it; its + teacher rows are read from the IPC records (:class:`SparseTeacherRowReader` + or :class:`DenseTeacherRowReader`) and its student values come from + :func:`distributed_selected_logprobs`. A mismatch chunk's student prefix + rows may sit on other CP ranks: each rank adds its share to one + fixed-length vector that is summed across CP with a SUM backward. + + The result is this rank's partial ``T**2 * (common + beta * mismatch) + / global_valid_chunks``; summed over CP ranks it equals the replicated + CSD KD term. Dense teachers follow the replicated support rules; sparse + teachers use each request's top-k with its realized label kept. + Returned metrics are CP sums. + """ + cfg = self.cfg + device = student_logits.device + accum_dtype = torch.float32 + temperature = float(cfg.temperature) + loss_fn_name = str(cfg.csd_loss_fn) + last_loss_fn_name = str(cfg.csd_last_pos_loss_fn or loss_fn_name) + mismatch_beta = float(cfg.csd_mismatch_loss_beta) + noise_filter_topk = int(cfg.csd_noise_filter_topk) + realized_only = bool(cfg.csd_realized_only_mismatch) + sparse = payload.transport == "sparse" + # Sparse owner-local transport always forces realized labels. + keep_realized = sparse + row_tile = int(cfg.csd_selected_row_chunk_size) + chunk_batch = int(cfg.csd_support_chunk_batch_size) + v_t = int(teacher_vocab_size) + tp_rank = ( + torch.distributed.get_rank(tp_group) + if tp_group is not None and torch.distributed.is_initialized() + else 0 + ) + cp_size = ( + torch.distributed.get_world_size(cp_group) + if cp_group is not None and torch.distributed.is_initialized() + else 1 + ) + local_vocab = int(student_logits.shape[-1]) + vocab_start = tp_rank * local_vocab + + def student_logp( + batch_rows: torch.Tensor, local_rows: torch.Tensor, token_ids: torch.Tensor + ) -> torch.Tensor: + return distributed_selected_logprobs( + student_logits, + batch_rows, + local_rows, + token_ids, + vocab_start_index=vocab_start, + vocab_end_index=vocab_start + local_vocab, + real_vocab_size=int(self.student_vocab_size), + temperature=temperature, + tp_group=tp_group, + row_chunk_size=row_tile, + ).to(accum_dtype) + + def as_long(values: list[int]) -> torch.Tensor: + return torch.tensor(values, device=device, dtype=torch.long) + + # A graph-connected zero: every rank's loss must reach the student logits + # (and the CP collectives below) even when it owns no chunk. + connected_zero = student_logits.reshape(-1)[0].to(accum_dtype) * 0.0 + + reader: Optional[SparseTeacherRowReader | DenseTeacherRowReader] = None + try: + setup_error: Optional[Exception] = None + # Collective-free setup: any failure here is agreed on below, before + # the first TP or CP collective. + try: + if align.student_global_positions is None: + raise ValueError("Owner-local CSD requires student positions.") + assert align.student_spans is not None + assert align.teacher_spans is not None + assert align.pair_valid is not None + assert align.student_input_ids is not None + assert align.teacher_input_ids is not None + sample_mask = to_local_if_dtensor(align.sample_mask).bool() + pair_valid = to_local_if_dtensor(align.pair_valid).bool() + batch_size = int(pair_valid.shape[0]) + if sample_mask.shape != (batch_size,): + raise ValueError( + "CSD sample_mask must have shape [B]: expected " + f"({batch_size},), got {tuple(sample_mask.shape)}." + ) + if len(payload.entries) != batch_size: + raise ValueError( + f"Teacher {i} IPC payload has {len(payload.entries)} " + f"samples for a microbatch of {batch_size}." + ) + pair_valid_np = ( + (pair_valid & sample_mask.unsqueeze(-1)).cpu().numpy().astype(bool) + ) + s_spans_np = align.student_spans.detach().cpu().numpy() + t_spans_np = align.teacher_spans.detach().cpu().numpy() + inp_s_np = align.student_input_ids.detach().cpu().numpy() + inp_t_np = align.teacher_input_ids.detach().cpu().numpy() + positions = align.student_global_positions.detach().cpu().tolist() + local_of_global = {int(g): j for j, g in enumerate(positions)} + student_seq_len = len(positions) * cp_size + + if sparse: + reader = SparseTeacherRowReader(payload.entries, device=device) + expected_membership = noise_filter_topk or int( + cfg.teacher_topk_ipc_k + ) + if ( + reader.k != int(cfg.teacher_topk_ipc_k) + or reader.temperature != temperature + or reader.real_vocab_size != v_t + or reader.membership_k != expected_membership + ): + raise ValueError( + f"Teacher {i} sparse IPC metadata (k={reader.k}, " + f"temperature={reader.temperature}, vocab=" + f"{reader.real_vocab_size}, membership_k=" + f"{reader.membership_k}) disagrees with the loss." + ) + else: + reader = DenseTeacherRowReader(payload.entries, device=device) + if reader.full_vocab_size < v_t: + raise ValueError( + f"Teacher {i} dense IPC vocabulary {reader.full_vocab_size} " + f"is narrower than the tokenizer vocabulary {v_t}." + ) + teacher_seq_len = reader.full_seq_len + common_s, common_t = self._get_common_indices_csd(device, i) + if common_s.numel() == 0: + raise ValueError( + f"Teacher {i}: the exact 1-to-1 common-vocab map is empty." + ) + + # Classify every chunk identically on every rank; keep the + # common chunks this rank owns and all mismatch chunks (their + # student prefixes may live here). + common_records: list[tuple[int, int, int, int, int]] = [] + mismatch_records: list[ + tuple[ + int, int, int, list[int], list[int], list[int], list[int], bool + ] + ] = [] + for b, k in zip(*np.nonzero(pair_valid_np)): + s_start, s_end = (int(x) for x in s_spans_np[b, k]) + t_start, t_end = (int(x) for x in t_spans_np[b, k]) + M, N = s_end - s_start, t_end - t_start + if M <= 0 or N <= 0: + continue + shift = ( + 1 if cfg.kl_chunk_shift and s_start > 0 and t_start > 0 else 0 + ) + s_labels = list(range(s_start, s_end)) + t_labels = list(range(t_start, t_end)) + s_pred = [pos - shift for pos in s_labels] + t_pred = [pos - shift for pos in t_labels] + if ( + s_pred[0] < 0 + or s_pred[-1] >= student_seq_len + or t_pred[0] < 0 + or t_pred[-1] >= teacher_seq_len + or s_labels[-1] >= inp_s_np.shape[1] + or t_labels[-1] >= inp_t_np.shape[1] + ): + continue + owned = s_pred[-1] in local_of_global + if M == 1 and N == 1: + if owned: + common_records.append( + ( + int(b), + local_of_global[s_pred[0]], + t_pred[0], + int(inp_s_np[b, s_labels[0]]), + int(inp_t_np[b, t_labels[0]]), + ) + ) + else: + mismatch_records.append( + (int(b), M, N, s_pred, s_labels, t_pred, t_labels, owned) + ) + + prefix_index = None + if not realized_only and any(r[-1] for r in mismatch_records): + prefix_index = self._ensure_bidir_prefix_support_index(device, i) + owned_mismatch: list[ + tuple[int, tuple[list[int], list[int], int, int]] + ] = [] + for record_idx, record in enumerate(mismatch_records): + if not record[-1]: + continue + support = self._owner_local_mismatch_support( + record, + inp_s_np=inp_s_np, + inp_t_np=inp_t_np, + prefix_index=prefix_index, + keep_realized=keep_realized, + ) + if support is not None: + owned_mismatch.append((record_idx, support)) + + # Sparse teacher rows are small (K values per request), so read + # every one this rank needs up front; this also validates the + # forced labels before any collective. + sparse_row_of: dict[tuple[int, int, int], int] = {} + sparse_rows: Optional[tuple[torch.Tensor, ...]] = None + if sparse: + assert isinstance(reader, SparseTeacherRowReader) + for b, _sl, t_pos, _rs, rt in common_records: + sparse_row_of.setdefault((b, t_pos, rt), len(sparse_row_of)) + for record_idx, _support in owned_mismatch: + b, _M, _N, _sp, _sl, t_pred, t_labels, _ow = mismatch_records[ + record_idx + ] + for t_pos, t_lab in zip(t_pred, t_labels): + key = (b, t_pos, int(inp_t_np[b, t_lab])) + sparse_row_of.setdefault(key, len(sparse_row_of)) + if sparse_row_of: + keys = list(sparse_row_of) + tiles = [ + reader.gather_rows( + as_long( + [key[0] for key in keys[start : start + row_tile]] + ), + as_long( + [key[1] for key in keys[start : start + row_tile]] + ), + as_long( + [key[2] for key in keys[start : start + row_tile]] + ), + ) + for start in range(0, len(keys), row_tile) + ] + sparse_rows = tuple( + torch.cat([tile[field] for tile in tiles]) + for field in range(4) + ) + except ( + AssertionError, + IndexError, + KeyError, + RuntimeError, + TypeError, + ValueError, + ) as exc: + setup_error = exc + self._raise_if_any_rank_failed( + setup_error, device=device, tp_group=tp_group, cp_group=cp_group + ) + + def sparse_label_logp( + keys: list[tuple[int, int, int]], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Teacher log-prob of each request's own label, and its membership.""" + assert sparse_rows is not None + rows = as_long([sparse_row_of[key] for key in keys]) + values, ids, log_z, in_topk = (field[rows] for field in sparse_rows) + match = ids.long() == as_long([key[2] for key in keys]).unsqueeze(-1) + raw = torch.where(match, values, torch.full_like(values, -torch.inf)) + return raw.max(dim=-1).values / temperature - log_z, in_topk + + def dense_rows( + batch_rows: list[int], teacher_rows: list[int] + ) -> tuple[torch.Tensor, torch.Tensor]: + """Real-vocabulary teacher logits ``[n, V_t]`` and their log Z.""" + assert isinstance(reader, DenseTeacherRowReader) + rows = reader.gather_rows(as_long(batch_rows), as_long(teacher_rows)) + rows = rows[:, :v_t] + return rows, torch.logsumexp(rows / temperature, dim=-1) + + def dense_label_in_topk( + rows: torch.Tensor, labels: torch.Tensor + ) -> torch.Tensor: + filter_k = min(noise_filter_topk, int(rows.shape[-1])) + topk_ids = torch.topk(rows, k=filter_k, dim=-1).indices + return (topk_ids == labels.unsqueeze(-1)).any(dim=-1) + + # ----- Common chunks owned by this rank ----------------------- + common_sum = connected_zero + common_count = 0 + common_matches = 0 + noise_filtered_common = 0 + teacher_to_common_student = ( + self._get_csd_teacher_to_common_student( + device, i, v_t, common_s, common_t + ) + if sparse + else None + ) + for start in range(0, len(common_records), chunk_batch): + batch = common_records[start : start + chunk_batch] + cb = as_long([r[0] for r in batch]) + cs = as_long([r[1] for r in batch]) + crs = as_long([r[3] for r in batch]) + crt = as_long([r[4] for r in batch]) + if sparse: + assert ( + sparse_rows is not None + and teacher_to_common_student is not None + ) + rows = as_long([sparse_row_of[(r[0], r[2], r[4])] for r in batch]) + topk_logits, topk_ids, teacher_log_z, in_topk = ( + field[rows] for field in sparse_rows + ) + mapped = teacher_to_common_student[topk_ids.long()] + # One of the K columns is reserved for the realized pair: + # keep the best K-1 other common alternatives by score and + # append the realized pair when it maps 1-to-1. + support = self._topk_mask_by_score( + (mapped >= 0) + & (mapped != crs.unsqueeze(-1)) + & (topk_ids.long() != crt.unsqueeze(-1)), + topk_logits.to(accum_dtype), + max(int(topk_ids.shape[-1]) - 1, 0), + ) + realized_raw = ( + torch.where( + topk_ids.long() == crt.unsqueeze(-1), + topk_logits, + torch.full_like(topk_logits, -torch.inf), + ) + .max(dim=-1) + .values + ) + realized_common = teacher_to_common_student[crt] == crs + queried = student_logp( + cb, + cs, + torch.cat([mapped.clamp_min(0), crs.unsqueeze(-1)], dim=-1), + ) + teacher_logp = torch.cat( + [ + topk_logits.to(accum_dtype) / temperature + - teacher_log_z.unsqueeze(-1), + (realized_raw / temperature - teacher_log_z).unsqueeze(-1), + ], + dim=-1, + ) + support = torch.cat( + [support, realized_common.unsqueeze(-1)], dim=-1 + ) + s_logp = torch.where(support, queried, -1.0e30) + t_logp = torch.where(support, teacher_logp, -1.0e30) + valid = support.any(dim=-1) + noise_ok = ( + in_topk if noise_filter_topk > 0 else torch.ones_like(valid) + ) + else: + rows, teacher_log_z = dense_rows( + [r[0] for r in batch], [r[2] for r in batch] + ) + support = None + s_logp = student_logp( + cb, cs, common_s.unsqueeze(0).expand(len(batch), -1) + ) + t_logp = rows[:, common_t] / temperature - teacher_log_z.unsqueeze( + -1 + ) + valid = torch.ones((len(batch),), device=device, dtype=torch.bool) + noise_ok = ( + dense_label_in_topk(rows, crt) + if noise_filter_topk > 0 + else torch.ones_like(valid) + ) + per_elem = self._partition_divergence( + self._append_rest_bucket_logp(s_logp), + self._append_rest_bucket_logp(t_logp), + loss_fn_name, + ) + if support is not None: + per_elem = per_elem * torch.cat( + [support.to(per_elem.dtype), torch.ones_like(per_elem[:, :1])], + dim=-1, + ) + keep = valid & noise_ok + common_sum = ( + common_sum + (per_elem.sum(dim=-1) * keep.to(accum_dtype)).sum() + ) + common_count += int(noise_ok.sum().item()) + noise_filtered_common += int((~noise_ok).sum().item()) + with torch.no_grad(): + common_matches += int( + ((s_logp.argmax(dim=-1) == t_logp.argmax(dim=-1)) & keep) + .sum() + .item() + ) + + # ----- Mismatch chunks ------------------------------------------ + mismatch_sum = connected_zero + mismatch_count = 0 + mismatch_matches = 0 + noise_filtered_mismatch = 0 + if mismatch_records: + # Each rank adds the realized student prefix log-probs of the + # rows it holds; the CP SUM gives every owner its complete prefix + # (realized-only chunks: the complete chain). + term_chunks: list[int] = [] + term_b: list[int] = [] + term_rows: list[int] = [] + term_labels: list[int] = [] + for record_idx, (b, _M, _N, s_pred, s_labels, *_rest) in enumerate( + mismatch_records + ): + terms = ( + zip(s_pred, s_labels) + if realized_only + else zip(s_pred[:-1], s_labels[:-1]) + ) + for pred_pos, label_pos in terms: + local_row = local_of_global.get(pred_pos) + if local_row is not None: + term_chunks.append(record_idx) + term_b.append(b) + term_rows.append(local_row) + term_labels.append(int(inp_s_np[b, label_pos])) + # Keep the vector in every rank's graph: a rank may own no final + # rows yet hold prefix rows, and must join the backward SUM. + local_terms = connected_zero + torch.zeros( + (len(mismatch_records),), device=device, dtype=accum_dtype + ) + if term_b: + local_terms = local_terms.index_add( + 0, + as_long(term_chunks), + student_logp( + as_long(term_b), as_long(term_rows), as_long(term_labels) + ), + ) + shared_terms = group_all_reduce_sum_with_grad_backward_sum( + local_terms, cp_group + ) + mismatch_sum = mismatch_sum + shared_terms.sum() * 0.0 + + for start in range(0, len(owned_mismatch), chunk_batch): + batch = owned_mismatch[start : start + chunk_batch] + records = [mismatch_records[idx] for idx, _ in batch] + n = len(batch) + student_prefix = shared_terms[as_long([idx for idx, _ in batch])] + # Teacher log-probs of every realized label in each chunk: + # the final row's label is needed only for noise filtering + # and realized-only chains. + requests = [ + (row, (r[0], t_pos, int(inp_t_np[r[0], t_lab])), j == r[2] - 1) + for row, r in enumerate(records) + for j, (t_pos, t_lab) in enumerate(zip(r[5], r[6])) + ] + owner = as_long([req[0] for req in requests]) + is_final = torch.tensor( + [req[2] for req in requests], device=device, dtype=torch.bool + ) + if sparse: + label_logp, label_ok = sparse_label_logp( + [req[1] for req in requests] + ) + else: + label_logp_parts, label_ok_parts = [], [] + for tile in range(0, len(requests), row_tile): + part = requests[tile : tile + row_tile] + rows, log_z = dense_rows( + [req[1][0] for req in part], [req[1][1] for req in part] + ) + labels = as_long([req[1][2] for req in part]) + label_logp_parts.append( + rows.gather(-1, labels.unsqueeze(-1)).squeeze(-1) + / temperature + - log_z + ) + label_ok_parts.append( + dense_label_in_topk(rows, labels) + if noise_filter_topk > 0 + else torch.ones_like(labels, dtype=torch.bool) + ) + label_logp = torch.cat(label_logp_parts) + label_ok = torch.cat(label_ok_parts) + chain_or_prefix_terms = ( + label_logp + if realized_only + else torch.where(is_final, 0.0, label_logp.to(accum_dtype)) + ) + teacher_prefix = torch.zeros( + (n,), device=device, dtype=accum_dtype + ).index_add(0, owner, chain_or_prefix_terms.to(accum_dtype)) + noise_ok = ( + torch.ones((n,), device=device, dtype=torch.int32) + .index_reduce(0, owner, label_ok.to(torch.int32), "amin") + .bool() + if noise_filter_topk > 0 + else torch.ones((n,), device=device, dtype=torch.bool) + ) + + if realized_only: + s_scores = student_prefix.unsqueeze(-1) + t_scores = teacher_prefix.unsqueeze(-1) + support = torch.ones((n, 1), device=device, dtype=torch.bool) + else: + width = max(len(sup[0]) for _, sup in batch) + s_ids = torch.zeros((n, width), device=device, dtype=torch.long) + t_ids = torch.zeros_like(s_ids) + candidate = torch.zeros( + (n, width), device=device, dtype=torch.bool + ) + realized_pair = torch.zeros_like(candidate) + for row, (_, (s_list, t_list, rs, rt)) in enumerate(batch): + count = len(s_list) + s_ids[row, :count] = as_long(s_list) + t_ids[row, :count] = as_long(t_list) + candidate[row, :count] = True + realized_pair[row, :count] = (s_ids[row, :count] == rs) & ( + t_ids[row, :count] == rt + ) + final_b = [r[0] for r in records] + final_t = [r[5][-1] for r in records] + if sparse: + assert sparse_rows is not None + final_rows = as_long( + [ + sparse_row_of[ + (b, t_pos, int(inp_t_np[b, r[6][-1]])) + ] + for b, t_pos, r in zip(final_b, final_t, records) + ] + ) + topk_logits, topk_ids, final_log_z, _ = ( + field[final_rows] for field in sparse_rows + ) + slot_match = t_ids.unsqueeze( + -1 + ) == topk_ids.long().unsqueeze(1) + found = slot_match.any(dim=-1) + t_raw = ( + torch.where( + slot_match, + topk_logits.unsqueeze(1), + torch.full_like( + topk_logits.unsqueeze(1), -torch.inf + ), + ) + .max(dim=-1) + .values + ) + t_final = t_raw.to( + accum_dtype + ) / temperature - final_log_z.unsqueeze(-1) + alternatives = self._topk_mask_by_score( + candidate & found & ~realized_pair, + torch.where(found, t_final, -1.0e30), + max(int(topk_ids.shape[-1]) - 1, 0), + ) + support = (realized_pair & found) | alternatives + else: + rows, final_log_z = dense_rows(final_b, final_t) + t_final = rows.gather( + -1, t_ids + ) / temperature - final_log_z.unsqueeze(-1) + support = candidate + final_s = as_long([local_of_global[r[3][-1]] for r in records]) + s_scores = student_logp(as_long(final_b), final_s, s_ids) + ( + student_prefix.unsqueeze(-1) + ) + t_scores = t_final + teacher_prefix.unsqueeze(-1) + s_scores = torch.where(support, s_scores, -1.0e30) + t_scores = torch.where(support, t_scores.to(accum_dtype), -1.0e30) + per_elem = self._partition_divergence( + self._append_rest_bucket_logp(s_scores), + self._append_rest_bucket_logp(t_scores), + last_loss_fn_name, + ) + per_elem = per_elem * torch.cat( + [support.to(per_elem.dtype), torch.ones_like(per_elem[:, :1])], + dim=-1, + ) + keep = support.any(dim=-1) & noise_ok + mismatch_sum = ( + mismatch_sum + + (per_elem.sum(dim=-1) * keep.to(accum_dtype)).sum() + ) + mismatch_count += int(keep.sum().item()) + noise_filtered_mismatch += int((~noise_ok).sum().item()) + with torch.no_grad(): + mismatch_matches += int( + ( + (s_scores.argmax(dim=-1) == t_scores.argmax(dim=-1)) + & keep + ) + .sum() + .item() + ) + finally: + if reader is not None: + reader.close() + + local_chunk_sum = common_sum + mismatch_beta * mismatch_sum + final_loss = ( + local_chunk_sum + / global_valid_chunks.to(device=device, dtype=accum_dtype).clamp(min=1.0) + * (temperature**2) + ) + stats = torch.tensor( + [ + float(common_sum.detach()), + float(mismatch_sum.detach()), + float(final_loss.detach()), + common_count, + mismatch_count, + common_matches, + mismatch_matches, + noise_filtered_common, + noise_filtered_mismatch, + ], + device=device, + dtype=torch.float64, + ) + if cp_size > 1: + torch.distributed.all_reduce(stats, group=cp_group) + ( + common_total, + mismatch_total, + kd_total, + common_total_count, + mismatch_total_count, + common_total_matches, + mismatch_total_matches, + noise_common_total, + noise_mismatch_total, + ) = stats.tolist() + mismatch_mean = mismatch_total / max(mismatch_total_count, 1.0) + metrics: dict[str, Any] = { + "loss": kd_total, + "train:loss": kd_total, + "kl_loss": kd_total, + "kl_common_per_chunk": common_total / max(common_total_count, 1.0), + "kl_partition_last_per_chunk": mismatch_mean, + "kl_mismatch_scaled_per_chunk": mismatch_mean * mismatch_beta, + "num_common_chunks": int(common_total_count), + "num_mismatch_chunks": int(mismatch_total_count), + "num_noise_filtered_common_chunks": int(noise_common_total), + "num_noise_filtered_mismatch_chunks": int(noise_mismatch_total), + "top1_acc_per_chunk": (common_total_matches + mismatch_total_matches) + / max(common_total_count + mismatch_total_count, 1.0), + "csd_mismatch_loss_beta": mismatch_beta, + } + if noise_filter_topk > 0: + metrics["csd_noise_filter_topk"] = noise_filter_topk + return final_loss, metrics diff --git a/nemo_rl/algorithms/loss/csd_packing.py b/nemo_rl/algorithms/loss/csd_packing.py new file mode 100644 index 00000000000..0aeb0cc6b5a --- /dev/null +++ b/nemo_rl/algorithms/loss/csd_packing.py @@ -0,0 +1,477 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sequence-packing adapter for cross-tokenizer distillation. + +The generic sequence-packing loss wrapper assumes every tensor with a second +dimension uses the student's token axis. CSD batches deliberately contain +three independent axes instead: student tokens, teacher tokens, and alignment +pair slots. This adapter restores each packed student row to its own effective +padded width and crops only fields on the student-token axis. +""" + +import math +from dataclasses import dataclass +from typing import Any, Callable, Optional, TypeVar + +import torch +import torch.distributed + +from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.model_utils import ( + _get_tokens_on_this_cp_rank, + allgather_cp_sharded_tensor, +) +from nemo_rl.utils.sequence_lengths import CpuIntTuple, to_cpu_int_tuple + +Tensor = TypeVar("Tensor", bound=torch.Tensor) + + +CSD_LOGICAL_METRICS_KEY = "__csd_logical_metrics__" + + +_STUDENT_SEQUENCE_KEYS = frozenset( + { + "input_ids", + "token_mask", + "kd_token_mask", + } +) +_STUDENT_ALIGNMENT_SEQUENCE_SUFFIXES = ( + "_student_chunk_id", + "_student_exact_partition_mask", +) + + +@dataclass(frozen=True) +class LockstepBinGeometry: + """The controller-owned geometry for one selected physical bin.""" + + bin_index: int + batch_item_ids: tuple[int, ...] + raw_cu_seqlens: tuple[int, ...] + padded_cu_seqlens: tuple[int, ...] + physical_tokens: int + + +def resolve_lockstep_bin_geometry( + data: BatchedDataDict[Any], + *, + input_lengths: torch.Tensor, +) -> LockstepBinGeometry | None: + """Resolve and validate the immutable plan geometry attached to ``data``. + + A physical-bin slice retains the complete rank-local bin-index tuple. The + exact ordered occurrence IDs identify which entry is being materialized. + Generic sequence packing has no lockstep metadata and returns ``None``. + Partially attached or drifting metadata fails closed. + """ + # Legacy callers and a few backend utility tests still pass a plain + # mapping. Absence of every attribute means generic (non-lockstep) + # packing; partially attached metadata continues to fail closed below. + plan = getattr(data, "lockstep_packing_plan", None) + side_id = getattr(data, "lockstep_side_id", None) + rank_bin_indices = getattr(data, "lockstep_bin_indices", None) + metadata = (plan, side_id, rank_bin_indices) + if all(value is None for value in metadata): + return None + if any(value is None for value in metadata): + raise ValueError( + "Lockstep packing metadata is incomplete: plan, side_id, and " + "rank-local bin indices must be attached together." + ) + assert plan is not None + assert side_id is not None + assert rank_bin_indices is not None + + if side_id not in plan.sides: + raise ValueError( + f"Lockstep plan batch_uid={plan.batch_uid} has no selected side " + f"{side_id!r}." + ) + side_plan = plan.sides[side_id] + normalized_rank_bins = tuple(int(value) for value in rank_bin_indices) + if normalized_rank_bins not in side_plan.rank_bin_indices: + raise ValueError( + f"Lockstep side {side_id!r} does not assign rank-local bins " + f"{normalized_rank_bins}." + ) + + if "batch_item_id" not in data: + raise ValueError( + "Plan-driven packing requires batch_item_id on every physical bin." + ) + item_id_values = data["batch_item_id"] + if isinstance(item_id_values, torch.Tensor): + if item_id_values.ndim != 1: + raise ValueError( + "batch_item_id must be one-dimensional, got shape " + f"{tuple(item_id_values.shape)}." + ) + batch_item_ids = tuple(int(value) for value in item_id_values.tolist()) + elif isinstance(item_id_values, list): + batch_item_ids = tuple(int(value) for value in item_id_values) + else: + raise TypeError( + "batch_item_id must be a one-dimensional tensor or list, got " + f"{type(item_id_values).__name__}." + ) + if len(set(batch_item_ids)) != len(batch_item_ids): + raise ValueError( + "batch_item_id values must be unique within a physical bin; got " + f"{batch_item_ids}." + ) + + matching_bins = [ + bin_index + for bin_index in normalized_rank_bins + if tuple(plan.bins[bin_index]) == batch_item_ids + ] + if len(matching_bins) != 1: + raise ValueError( + "Physical-bin batch_item_id order does not identify exactly one " + f"assigned lockstep bin: ids={batch_item_ids}, assigned_bins=" + f"{normalized_rank_bins}." + ) + bin_index = matching_bins[0] + + if input_lengths.ndim != 1 or int(input_lengths.numel()) != len(batch_item_ids): + raise ValueError( + "input_lengths must contain one value per planned physical-bin item; " + f"got shape {tuple(input_lengths.shape)} for {len(batch_item_ids)} IDs." + ) + actual_lengths = tuple(int(value) for value in input_lengths.tolist()) + raw_cu_seqlens = tuple(side_plan.raw_cu_seqlens_by_bin[bin_index]) + planned_lengths = tuple( + right - left for left, right in zip(raw_cu_seqlens, raw_cu_seqlens[1:]) + ) + if actual_lengths != planned_lengths: + raise ValueError( + f"Lockstep side {side_id!r} bin {bin_index} input lengths drifted: " + f"actual={actual_lengths}, planned={planned_lengths}." + ) + + padded_cu_seqlens = tuple(side_plan.padded_cu_seqlens_by_bin[bin_index]) + physical_tokens = int(side_plan.physical_tokens_by_bin[bin_index]) + if padded_cu_seqlens[-1] != physical_tokens: + raise ValueError( + f"Lockstep side {side_id!r} bin {bin_index} padded boundary ends at " + f"{padded_cu_seqlens[-1]}, not physical size {physical_tokens}." + ) + return LockstepBinGeometry( + bin_index=bin_index, + batch_item_ids=batch_item_ids, + raw_cu_seqlens=raw_cu_seqlens, + padded_cu_seqlens=padded_cu_seqlens, + physical_tokens=physical_tokens, + ) + + +def _context_parallel_size( + context_parallel_group: Optional[torch.distributed.ProcessGroup], +) -> int: + if context_parallel_group is None: + return 1 + return torch.distributed.get_world_size(context_parallel_group) + + +def _packed_row_geometry( + *, + sequence_index: int, + cu_seqlens_q: torch.Tensor | CpuIntTuple, + cu_seqlens_q_padded: torch.Tensor | CpuIntTuple, +) -> tuple[int, int, int]: + """Return ``(raw_length, padded_start, padded_length)`` for one row.""" + for boundaries in (cu_seqlens_q, cu_seqlens_q_padded): + if isinstance(boundaries, torch.Tensor) and boundaries.ndim != 1: + raise ValueError("Packed CSD cu_seqlens must be one-dimensional.") + cu_seqlens_q = to_cpu_int_tuple(cu_seqlens_q) + cu_seqlens_q_padded = to_cpu_int_tuple(cu_seqlens_q_padded) + num_sequences = len(cu_seqlens_q) - 1 + if len(cu_seqlens_q_padded) != num_sequences + 1: + raise ValueError( + "Raw and padded cu_seqlens must describe the same number of sequences." + ) + if sequence_index < 0 or sequence_index >= num_sequences: + raise IndexError( + f"sequence_index={sequence_index} is outside [0, {num_sequences})." + ) + + raw_start = cu_seqlens_q[sequence_index] + raw_end = cu_seqlens_q[sequence_index + 1] + padded_start = cu_seqlens_q_padded[sequence_index] + padded_end = cu_seqlens_q_padded[sequence_index + 1] + raw_length = raw_end - raw_start + padded_length = padded_end - padded_start + if ( + raw_start < 0 + or padded_start < 0 + or raw_length < 0 + or padded_length < raw_length + ): + raise ValueError( + "Invalid packed CSD geometry for sequence " + f"{sequence_index}: raw=({raw_start}, {raw_end}), " + f"padded=({padded_start}, {padded_end})." + ) + return raw_length, padded_start, padded_length + + +def _is_student_sequence_key(key: str) -> bool: + """Return whether ``key`` has the student's token axis in dimension one.""" + return key in _STUDENT_SEQUENCE_KEYS or ( + key.startswith("alignment_") + and key.endswith(_STUDENT_ALIGNMENT_SEQUENCE_SUFFIXES) + ) + + +def _crop_student_sequence_tensors( + data: BatchedDataDict[Any], + *, + sequence_length: int, +) -> BatchedDataDict[Any]: + """Crop only known student-token tensors to one row's effective width.""" + for key, value in tuple(data.items()): + if not _is_student_sequence_key(key): + continue + if not isinstance(value, torch.Tensor): + raise TypeError( + f"Packed CSD student sequence field {key!r} must be a " + f"tensor, got {type(value).__name__}." + ) + if value.ndim < 2: + raise ValueError( + f"Packed CSD student sequence field {key!r} must have a " + f"sequence axis in dimension one, got shape {tuple(value.shape)}." + ) + if int(value.shape[1]) < sequence_length: + raise ValueError( + f"Packed CSD student sequence field {key!r} has width " + f"{value.shape[1]}, smaller than row effective width " + f"{sequence_length}." + ) + if int(value.shape[1]) > sequence_length: + data[key] = value.narrow(1, 0, sequence_length) + return data + + +def restore_packed_student_logits( + packed_logits: torch.Tensor, + *, + sequence_index: int, + cu_seqlens_q: torch.Tensor | CpuIntTuple, + cu_seqlens_q_padded: torch.Tensor | CpuIntTuple, + context_parallel_group: Optional[torch.distributed.ProcessGroup] = None, +) -> torch.Tensor: + """Restore one logical student's logits without touching other data axes. + + ``packed_logits`` follows the backend's packed THD layout. With Megatron + context parallelism, every physical sequence is independently sharded in + head/tail order. The returned tensor follows the normal, unpacked backend + contract: ``[1, S_row, V]`` for CP1, or this CP rank's head/tail shard + ``[1, S_row / CP, V]`` for CP>1. ``S_row`` is the row's own padded-boundary + delta, not the batch-wide student rectangle. Padding positions are zero and + remain masked by the cropped logical row's token mask. + """ + if packed_logits.ndim != 3 or packed_logits.shape[0] != 1: + raise ValueError( + "Packed CSD logits must have shape [1, T, V], got " + f"{tuple(packed_logits.shape)}." + ) + raw_length, padded_start, padded_length = _packed_row_geometry( + sequence_index=sequence_index, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_q_padded=cu_seqlens_q_padded, + ) + + cp_size = _context_parallel_size(context_parallel_group) + if padded_start % cp_size != 0 or padded_length % cp_size != 0: + raise ValueError( + "Packed padded boundaries must be divisible by context-parallel " + f"size {cp_size}; got start={padded_start}, length={padded_length}." + ) + local_start = padded_start // cp_size + local_length = padded_length // cp_size + if local_start + local_length > packed_logits.shape[1]: + raise ValueError( + "Packed logit slice exceeds the local sequence dimension: " + f"start={local_start}, length={local_length}, " + f"available={packed_logits.shape[1]}." + ) + + physical_local = packed_logits.narrow(1, local_start, local_length) + if cp_size > 1: + assert context_parallel_group is not None + physical_full = allgather_cp_sharded_tensor( + physical_local, context_parallel_group, seq_dim=1 + ) + else: + physical_full = physical_local + valid_logits = physical_full.narrow(1, 0, raw_length) + if raw_length < padded_length: + valid_logits = torch.nn.functional.pad( + valid_logits, (0, 0, 0, padded_length - raw_length) + ) + + if cp_size == 1: + return valid_logits + if padded_length % (2 * cp_size) != 0: + raise ValueError( + "The effective student sequence width must be divisible by 2 * " + f"context-parallel size; got width={padded_length}, " + f"cp_size={cp_size}." + ) + assert context_parallel_group is not None + cp_rank = torch.distributed.get_rank(context_parallel_group) + return _get_tokens_on_this_cp_rank( + valid_logits, + cp_rank=cp_rank, + cp_size=cp_size, + seq_dim=1, + ).contiguous() + + +class CSDSequencePackingLossWrapper: + """Run cross-tokenizer distillation at logical microbatch size one. + + Unlike :class:`SequencePackingLossWrapper`, this wrapper distinguishes the + independent student-token, teacher-token, and pair-slot axes. A sliced + logical row restores and crops only its student-token rectangle to that + row's padded boundary delta. Teacher, pair, region, and IPC payloads retain + their original per-row shapes and contents. + """ + + def __init__( + self, + loss_fn: LossFunction, + prepare_fn: Callable[..., Any], + cu_seqlens_q: Tensor | CpuIntTuple, + cu_seqlens_q_padded: Optional[Tensor | CpuIntTuple] = None, + vocab_parallel_rank: Optional[int] = None, + vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, + context_parallel_group: Optional[torch.distributed.ProcessGroup] = None, + ) -> None: + self.loss_fn = loss_fn + self.prepare_fn = prepare_fn + self.cu_seqlens_q = to_cpu_int_tuple(cu_seqlens_q) + self.cu_seqlens_q_padded = to_cpu_int_tuple( + cu_seqlens_q if cu_seqlens_q_padded is None else cu_seqlens_q_padded + ) + self.vocab_parallel_rank = vocab_parallel_rank + self.vocab_parallel_group = vocab_parallel_group + self.context_parallel_group = context_parallel_group + + def __call__( + self, + next_token_logits: Tensor, + data: BatchedDataDict[Any], + global_valid_seqs: Tensor, + global_valid_toks: Tensor, + global_valid_kd_toks: Tensor | None = None, + global_valid_chunks_by_idx: dict[int, torch.Tensor] | None = None, + ) -> tuple[Tensor, dict[str, Any]]: + num_sequences = len(self.cu_seqlens_q) - 1 + if data.size != num_sequences: + raise ValueError( + "Packed CSD metadata must contain one boundary per logical " + f"row: data.size={data.size}, sequences={num_sequences}." + ) + loss_accum: Any = 0 + metrics_accum: dict[str, Any] = {} + logical_metrics: list[dict[str, Any]] = [] + for sequence_index in range(num_sequences): + _, _, effective_sequence_length = _packed_row_geometry( + sequence_index=sequence_index, + cu_seqlens_q=self.cu_seqlens_q, + cu_seqlens_q_padded=self.cu_seqlens_q_padded, + ) + logical_logits = restore_packed_student_logits( + next_token_logits, + sequence_index=sequence_index, + cu_seqlens_q=self.cu_seqlens_q, + cu_seqlens_q_padded=self.cu_seqlens_q_padded, + context_parallel_group=self.context_parallel_group, + ) + sequence_data = data.slice(sequence_index, sequence_index + 1) + sequence_data = _crop_student_sequence_tensors( + sequence_data, + sequence_length=effective_sequence_length, + ) + loss_input, sequence_data = self.prepare_fn( + logits=logical_logits, + data=sequence_data, + loss_fn=self.loss_fn, + vocab_parallel_rank=self.vocab_parallel_rank, + vocab_parallel_group=self.vocab_parallel_group, + context_parallel_group=self.context_parallel_group, + ) + extra_loss_kwargs: dict[str, Any] = {} + if global_valid_chunks_by_idx: + extra_loss_kwargs["global_valid_chunks_by_idx"] = ( + global_valid_chunks_by_idx + ) + loss, metrics = self.loss_fn( + data=sequence_data, + global_valid_seqs=global_valid_seqs, + global_valid_toks=global_valid_toks, + global_valid_kd_toks=global_valid_kd_toks, + **extra_loss_kwargs, + **loss_input, + ) + loss_accum += loss + self._accumulate_metrics(metrics_accum, metrics) + + # The worker's established aggregation treats every metric dict as + # one logical MBS1 cohort. Preserve that cardinality instead of + # turning an uneven physical bin into one intensive-metric record. + # Masked rows still execute the loss (and any collectives), but + # match DTensor's unpacked behavior by contributing no metric row. + num_valid_samples = metrics.get("num_valid_samples", 1) + if isinstance(num_valid_samples, torch.Tensor): + num_valid_samples = num_valid_samples.item() + if num_valid_samples > 0: + logical_metrics.append(dict(metrics)) + + metrics_accum[CSD_LOGICAL_METRICS_KEY] = logical_metrics + return loss_accum, metrics_accum + + @staticmethod + def _accumulate_metrics( + accumulated: dict[str, Any], metrics: dict[str, Any] + ) -> None: + """Match the established packing wrapper's metric reductions.""" + minimum_metrics = {"probs_ratio_min", "probs_ratio_clamped_min"} + maximum_metrics = {"probs_ratio_max", "probs_ratio_clamped_max"} + for key, value in metrics.items(): + if key not in accumulated: + if key in minimum_metrics: + accumulated[key] = float("inf") + elif key in maximum_metrics: + accumulated[key] = float("-inf") + else: + accumulated[key] = 0 + scalar_or_value = ( + value.item() + if isinstance(value, torch.Tensor) and value.ndim == 0 + else value + ) + if key in minimum_metrics: + if not math.isinf(scalar_or_value): + accumulated[key] = min(accumulated[key], scalar_or_value) + elif key in maximum_metrics: + if not math.isinf(scalar_or_value): + accumulated[key] = max(accumulated[key], scalar_or_value) + else: + accumulated[key] += scalar_or_value diff --git a/nemo_rl/algorithms/loss/loss_functions.py b/nemo_rl/algorithms/loss/loss_functions.py index 9330471a04c..f13de33c7d2 100755 --- a/nemo_rl/algorithms/loss/loss_functions.py +++ b/nemo_rl/algorithms/loss/loss_functions.py @@ -1907,7 +1907,7 @@ def __call__( metrics = { "loss": float(kl_loss.item()) if kl_loss.ndim == 0 else kl_loss, - "num_valid_samples": data["input_ids"].shape[0], + "num_valid_samples": data["sample_mask"].sum().item(), } return kl_loss, metrics @@ -2042,7 +2042,7 @@ def __call__( "values_max": values_max, "returns_sq_mean": returns_sq_mean, "residual_sq_mean": residual_sq_mean, - "num_valid_samples": int(values.shape[0]), + "num_valid_samples": sample_mask.sum().item(), } return loss, metrics diff --git a/nemo_rl/algorithms/loss/loss_input.py b/nemo_rl/algorithms/loss/loss_input.py index 9dd4a96d785..ef5728fefd9 100644 --- a/nemo_rl/algorithms/loss/loss_input.py +++ b/nemo_rl/algorithms/loss/loss_input.py @@ -22,6 +22,7 @@ TrainingSamplingParams, need_top_k_or_top_p_filtering, ) +from nemo_rl.algorithms.loss.csd import CSDLossFn from nemo_rl.algorithms.loss.interfaces import LossFunction, LossInputType from nemo_rl.algorithms.loss.utils import ( _pack_input_ids, @@ -29,6 +30,8 @@ reconstruct_opd_full_teacher_logits, ) from nemo_rl.algorithms.x_token.loss_utils import ( + prepare_csd_cross_tokenizer_loss_input, + prepare_csd_owner_local_loss_input, prepare_xtoken_cross_tokenizer_loss_input, ) from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -325,34 +328,101 @@ def prepare_loss_input( } elif loss_fn.input_type == LossInputType.DISTILLATION_CROSS_TOKENIZER: - # Rebuild each teacher's full-vocab logits from its per-rank CUDA IPC - # handles and do the shared CP-resolution the loss needs; the loss fn - # does the per-teacher projection / chunk-average / KL reductions and - # aggregates them by ``kd_loss_mode``. ``projection_matrix_paths`` drives - # the teacher count and which teachers are same-tokenizer (``None``). The - # TP group is derived from the student logits' own device mesh. - ( - student_logits_contig, - teacher_full_logits_by_idx, - aligns_by_idx, - tp_group, - cp_group, - ) = prepare_xtoken_cross_tokenizer_loss_input( - logits, - data, - projection_matrix_paths=loss_fn.projection_matrix_paths, - vocab_parallel_group=vocab_parallel_group, - context_parallel_group=context_parallel_group, - cp_sharder=cp_sharder, - ) - loss_input = { - "logits": logits, - "student_logits_contig": student_logits_contig, - "teacher_full_logits_by_idx": teacher_full_logits_by_idx, - "aligns_by_idx": aligns_by_idx, - "tp_group": tp_group, - "cp_group": cp_group, - } + if isinstance(loss_fn, CSDLossFn) and loss_fn.owner_local: + # Megatron student logits stay in their CP shard; cross-tokenizer + # teacher rows are read by the loss from their IPC records. + ( + student_logits_contig, + teacher_full_logits_by_idx, + teacher_row_payloads_by_idx, + aligns_by_idx, + tp_group, + cp_group, + dp_cp_group, + ) = prepare_csd_owner_local_loss_input( + logits, + data, + cross_tokenizer_teacher_mask=[ + not loss_fn._teacher_is_same_vocab(i) + for i in range(loss_fn.num_teachers) + ], + vocab_parallel_group=vocab_parallel_group, + context_parallel_group=context_parallel_group, + ) + loss_input = { + "logits": logits, + "student_logits_contig": student_logits_contig, + "teacher_full_logits_by_idx": teacher_full_logits_by_idx, + "teacher_row_payloads_by_idx": teacher_row_payloads_by_idx, + "aligns_by_idx": aligns_by_idx, + "tp_group": tp_group, + "cp_group": cp_group, + "dp_cp_group": dp_cp_group, + } + elif isinstance(loss_fn, CSDLossFn): + ( + student_logits_contig, + teacher_full_logits_by_idx, + teacher_sparse_logits_by_idx, + aligns_by_idx, + dense_reconstruction_fallbacks_by_idx, + tp_group, + cp_group, + dp_cp_group, + ) = prepare_csd_cross_tokenizer_loss_input( + logits, + data, + projection_matrix_paths=loss_fn.projection_matrix_paths, + cross_tokenizer_teacher_mask=[ + not loss_fn._teacher_is_same_vocab(i) + for i in range(loss_fn.num_teachers) + ], + vocab_parallel_group=vocab_parallel_group, + context_parallel_group=context_parallel_group, + cp_sharder=cp_sharder, + ) + loss_input = { + "logits": logits, + "student_logits_contig": student_logits_contig, + "teacher_full_logits_by_idx": teacher_full_logits_by_idx, + "teacher_sparse_logits_by_idx": teacher_sparse_logits_by_idx, + "aligns_by_idx": aligns_by_idx, + "dense_reconstruction_fallbacks_by_idx": dense_reconstruction_fallbacks_by_idx, + "tp_group": tp_group, + "cp_group": cp_group, + "dp_cp_group": dp_cp_group, + # Megatron (no model-owned CP sharder) divides the loss by CP. + "cp_normalized_by_caller": cp_sharder is None, + } + else: + # Rebuild each teacher's full-vocab logits from its per-rank CUDA IPC + # handles and do the shared CP-resolution the loss needs; the loss fn + # does the per-teacher projection / chunk-average / KL reductions and + # aggregates them by ``kd_loss_mode``. ``projection_matrix_paths`` drives + # the teacher count and which teachers are same-tokenizer (``None``). The + # TP group is derived from the student logits' own device mesh. + ( + student_logits_contig, + teacher_full_logits_by_idx, + aligns_by_idx, + tp_group, + cp_group, + ) = prepare_xtoken_cross_tokenizer_loss_input( + logits, + data, + projection_matrix_paths=loss_fn.projection_matrix_paths, + vocab_parallel_group=vocab_parallel_group, + context_parallel_group=context_parallel_group, + cp_sharder=cp_sharder, + ) + loss_input = { + "logits": logits, + "student_logits_contig": student_logits_contig, + "teacher_full_logits_by_idx": teacher_full_logits_by_idx, + "aligns_by_idx": aligns_by_idx, + "tp_group": tp_group, + "cp_group": cp_group, + } if cp_sharder is not None: next_token_logprobs = get_cp_sharded_next_token_logprobs( logits, diff --git a/nemo_rl/algorithms/loss/wrapper.py b/nemo_rl/algorithms/loss/wrapper.py index 79fa197fa10..7c3d23ac751 100644 --- a/nemo_rl/algorithms/loss/wrapper.py +++ b/nemo_rl/algorithms/loss/wrapper.py @@ -89,6 +89,7 @@ def __call__( data: BatchedDataDict[Any], global_valid_seqs: Tensor | None, global_valid_toks: Tensor | None, + global_valid_chunks_by_idx: dict[int, torch.Tensor] | None = None, ) -> tuple[Tensor, dict[str, Any]]: """Wraps a loss function to handle sequence packing by doing one sequence at a time to avoid excessive padding.""" unpadded_cu_seqlens = self.cu_seqlens_q @@ -152,10 +153,16 @@ def __call__( ) # call loss function + extra_loss_kwargs: dict[str, Any] = {} + if global_valid_chunks_by_idx: + extra_loss_kwargs["global_valid_chunks_by_idx"] = ( + global_valid_chunks_by_idx + ) loss, metrics = self.loss_fn( data=unpadded_seq_data, global_valid_seqs=global_valid_seqs, global_valid_toks=global_valid_toks, + **extra_loss_kwargs, **loss_input, ) @@ -227,6 +234,7 @@ def __call__( data: BatchedDataDict[Any], global_valid_seqs: Tensor | None, global_valid_toks: Tensor | None, + global_valid_chunks_by_idx: dict[int, torch.Tensor] | None = None, ) -> tuple[Tensor, dict[str, Any]]: """Compute loss for all packed sequences in one forward pass.""" loss_input, prepared_data = self.prepare_fn( @@ -240,10 +248,14 @@ def __call__( context_parallel_group=self.context_parallel_group, ) + extra_loss_kwargs: dict[str, Any] = {} + if global_valid_chunks_by_idx: + extra_loss_kwargs["global_valid_chunks_by_idx"] = global_valid_chunks_by_idx return self.loss_fn( data=prepared_data, global_valid_seqs=global_valid_seqs, global_valid_toks=global_valid_toks, + **extra_loss_kwargs, **loss_input, ) @@ -425,6 +437,8 @@ def wrap_loss_fn_with_input_preparation( vocab_parallel_rank: Optional[int] = None, vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, context_parallel_group: Optional[torch.distributed.ProcessGroup] = None, + global_valid_kd_toks: Tensor | None = None, + global_valid_chunks_by_idx: dict[int, torch.Tensor] | None = None, ) -> tuple[Tensor, dict[str, Any]]: """Wraps a loss function to handle input preparation for megatron policy worker.""" # prepare loss input @@ -438,10 +452,16 @@ def wrap_loss_fn_with_input_preparation( ) # call loss function + extra_loss_kwargs: dict[str, Any] = {} + if global_valid_chunks_by_idx: + extra_loss_kwargs["global_valid_chunks_by_idx"] = global_valid_chunks_by_idx + if global_valid_kd_toks is not None: + extra_loss_kwargs["global_valid_kd_toks"] = global_valid_kd_toks loss, loss_metrics = loss_fn( data=data, global_valid_seqs=global_valid_seqs, global_valid_toks=global_valid_toks, + **extra_loss_kwargs, **loss_input, ) diff --git a/nemo_rl/algorithms/x_token/loss_utils.py b/nemo_rl/algorithms/x_token/loss_utils.py index 72dfbb10bc4..d1e1dac5ef9 100644 --- a/nemo_rl/algorithms/x_token/loss_utils.py +++ b/nemo_rl/algorithms/x_token/loss_utils.py @@ -24,9 +24,12 @@ :func:`nemo_rl.distributed.model_utils.group_all_reduce_sum` for the global valid-chunk denominator. - Teacher-logit IPC: :func:`rebuild_teacher_full_logits_from_ipc`, + :func:`rebuild_teacher_sparse_logits_from_ipc`, :func:`assemble_teacher_logits_from_shards`, :func:`collect_overlapping_teacher_shards` reassemble full-vocab teacher - logits from per-rank shards across heterogeneous TP/CP. + logits from per-rank shards across heterogeneous TP/CP; + :class:`SparseTeacherRowReader` / :class:`DenseTeacherRowReader` read only + requested teacher rows. - Projection: :func:`parse_projection_file`, the :func:`get_sparse_projection_matrix` / :func:`get_topk_projection` process-local caches, :func:`slice_sparse_projection_rows`, and @@ -39,15 +42,22 @@ from __future__ import annotations +import math import os -from dataclasses import dataclass +from bisect import bisect_right +from collections.abc import Sequence +from dataclasses import dataclass, fields from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Tuple, Union import torch from torch.distributed.tensor import DTensor from nemo_rl.algorithms.utils import masked_mean +from nemo_rl.algorithms.x_token.token_aligner import AlignmentBatch from nemo_rl.distributed.model_utils import ( + _get_tokens_on_this_cp_rank, + allgather_cp_contiguous_tensor, + cp_load_balanced_to_contiguous, cp_shift_next, get_logprobs_from_vocab_parallel_logits, group_all_reduce_sum_with_grad, @@ -63,6 +73,25 @@ ) +SparseTeacherLogits = tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + Optional[torch.Tensor], +] + + +def alignment_from_flat_batch(data: Mapping[str, Any]) -> AlignmentBatch: + """Rebuild :class:`AlignmentBatch` from the flat ``alignment_*`` keys. + + The field set is driven off :class:`AlignmentBatch` so the helper + can't drift from the schema. + """ + return AlignmentBatch( + **{f.name: data[f"alignment_{f.name}"] for f in fields(AlignmentBatch)} + ) + + class Fp32SparseMM(torch.autograd.Function): """FP32 ``M.t() @ dense`` (sparse-dense matmul) ignoring surrounding autocast. @@ -504,6 +533,51 @@ class LocalizedAlignment: # next-token-accuracy metric and the same-tokenizer KD path. student_input_ids: Optional[torch.Tensor] = None student_token_mask: Optional[torch.Tensor] = None + # KD semantic targets may be narrower than the SFT/CE role mask in chat + # mode (assistant content plus an explicitly identified EOT token). Text + # mode aliases this to ``student_token_mask``. + student_kd_token_mask: Optional[torch.Tensor] = None + # Filled post-construction for the CSD (CSD) path: + # this CP rank's contiguous teacher input ids, per-chunk contiguous position + # spans ``[B, max_pairs, 2]`` derived from the unshifted chunk ids, and the + # per-sample chunk count ``[B]``. Left ``None`` for a same-tokenizer teacher. + teacher_input_ids: Optional[torch.Tensor] = None + student_spans: Optional[torch.Tensor] = None + teacher_spans: Optional[torch.Tensor] = None + num_chunks: Optional[torch.Tensor] = None + # Owner-local CSD only: the global student position of each local logit row + # (the load-balanced CP layout). Ids, masks and spans then stay global. + student_global_positions: Optional[torch.Tensor] = None + + +@dataclass(frozen=True) +class TeacherRowPayload: + """One cross-tokenizer teacher's per-sample IPC records, read row by row. + + ``transport`` is ``"sparse"`` (records from ``get_sparse_logits_ipc``, read + with :class:`SparseTeacherRowReader`) or ``"dense"`` (full-logit records, + read with :class:`DenseTeacherRowReader`). + """ + + transport: str + entries: list[dict[str, Any]] + + +def cp_load_balanced_global_positions( + local_seq_len: int, + cp_group: Optional[torch.distributed.ProcessGroup], + *, + device: torch.device, +) -> torch.Tensor: + """Global positions ``[local_seq_len]`` of this CP rank's load-balanced rows.""" + cp_size = ( + torch.distributed.get_world_size(cp_group) + if cp_group is not None and torch.distributed.is_initialized() + else 1 + ) + cp_rank = torch.distributed.get_rank(cp_group) if cp_size > 1 else 0 + full = torch.arange(local_seq_len * cp_size, device=device, dtype=torch.long) + return _get_tokens_on_this_cp_rank(full, cp_rank, cp_size, seq_dim=0) def localize_alignment( @@ -541,6 +615,7 @@ def student_next_token_ce( logits: torch.Tensor, *, input_ids: torch.Tensor, + seq_index: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Per-token next-token cross-entropy ``[B, T-1]`` on the student. @@ -549,7 +624,11 @@ def student_next_token_ce( last predictor) matches the convention the KL terms use. """ if isinstance(logits, DTensor): - next_token_logprobs = get_logprobs_from_vocab_parallel_logits(logits, input_ids) + next_token_logprobs = get_logprobs_from_vocab_parallel_logits( + logits, + input_ids, + **({"seq_index": seq_index} if seq_index is not None else {}), + ) return -next_token_logprobs shift_logits = logits[:, :-1].contiguous() shift_labels = input_ids[:, 1:].contiguous() @@ -625,12 +704,24 @@ def collect_overlapping_teacher_shards( student_seq_start = student_cp_rank * full_seq_len // student_cp_size student_seq_end = (student_cp_rank + 1) * full_seq_len // student_cp_size + from nemo_rl.models.policy.utils import validate_compact_teacher_ipc_handle + matches: list[tuple[dict[str, Any], slice, slice, slice, slice]] = [] for handle in teacher_shards: teacher_vocab_start = int(handle["vocab_start_index"]) teacher_vocab_end = int(handle["vocab_end_index"]) teacher_seq_start = int(handle["global_seq_start"]) - teacher_seq_end = teacher_seq_start + int(handle["actual_shape"][0]) + compact_geometry = validate_compact_teacher_ipc_handle(handle) + stored_seq_len = ( + compact_geometry[1] + if compact_geometry is not None + else int(handle["actual_shape"][0]) + ) + # A compact shard's omitted suffix is logically present but known zero. + # Plan copies only for the physically stored prefix; ``dest`` is + # zero-initialized by the caller and therefore reconstructs the exact + # old dense rectangle. + teacher_seq_end = teacher_seq_start + stored_seq_len overlap_seq_start = max(student_seq_start, teacher_seq_start) overlap_seq_end = min(student_seq_end, teacher_seq_end) @@ -651,6 +742,52 @@ def collect_overlapping_teacher_shards( return matches +def _rebuild_compact_teacher_ipc_storage( + handle: Mapping[str, Any], device: int +) -> tuple[torch.Tensor, tuple[int, int, int, int]]: + """Map and validate one compact producer slab without copying it.""" + from nemo_rl.models.policy.utils import ( + rebuild_cuda_tensor_from_ipc, + validate_compact_teacher_ipc_handle, + ) + + compact_geometry = validate_compact_teacher_ipc_handle(handle) + if compact_geometry is None: + raise ValueError("Expected a compact dense teacher IPC handle.") + _token_offset, stored_seq_len, _used_tokens, _local_vocab_size = compact_geometry + if stored_seq_len <= 0: + raise ValueError("A zero-length compact IPC row has no storage to rebuild.") + src_full = rebuild_cuda_tensor_from_ipc(handle["payload_ipc"], device).detach() + expected_storage_shape = tuple(int(value) for value in handle["storage_shape"]) + if tuple(src_full.shape) != expected_storage_shape or src_full.ndim != 2: + raise ValueError( + "Compact dense teacher IPC rebuilt an unexpected storage shape: " + f"actual={tuple(src_full.shape)}, expected={expected_storage_shape}." + ) + if not src_full.is_contiguous(): + raise ValueError("Compact dense teacher IPC storage must be contiguous.") + if src_full.dtype != handle["dtype"]: + raise ValueError( + "Compact dense teacher IPC rebuilt an unexpected dtype: " + f"actual={src_full.dtype}, expected={handle['dtype']}." + ) + return src_full, compact_geometry + + +def _rebuild_teacher_ipc_row(handle: Mapping[str, Any], device: int) -> torch.Tensor: + """Return the physically stored ``[T_stored, V_local]`` row view.""" + from nemo_rl.models.policy.utils import ( + rebuild_cuda_tensor_from_ipc, + validate_compact_teacher_ipc_handle, + ) + + if validate_compact_teacher_ipc_handle(handle) is not None: + src_full, _ = _rebuild_compact_teacher_ipc_storage(handle, device) + else: + src_full = rebuild_cuda_tensor_from_ipc(handle["payload_ipc"], device).detach() + return _teacher_ipc_row_view(handle, src_full) + + def assemble_teacher_logits_from_shards( teacher_shards: list[dict[str, Any]], student_cp_rank: int, @@ -662,8 +799,6 @@ def assemble_teacher_logits_from_shards( ``device`` is a CUDA device index (matches :func:`rebuild_cuda_tensor_from_ipc`'s ``device_id`` signature). """ - from nemo_rl.models.policy.utils import rebuild_cuda_tensor_from_ipc - if not teacher_shards: raise ValueError("teacher_shards must be non-empty") full_seq_len = int(teacher_shards[0]["full_seq_len"]) @@ -688,14 +823,7 @@ def assemble_teacher_logits_from_shards( full_seq_len=full_seq_len, ) for handle, src_seq, src_vocab, dest_seq, dest_vocab in matches: - # Producer's IPC payload is the full contiguous storage - # [N_microbatches, B_mb, T_t_local, V_t_local]; index the slot - # then the sample row, then apply the seq/vocab overlap slices. - src_full = rebuild_cuda_tensor_from_ipc(handle["payload_ipc"], device).detach() - buf_idx = int(handle["buf_idx"]) - sample_idx = int(handle["sample_index_in_buf"]) - local_seq_t, local_vocab_t = handle["actual_shape"] - src = src_full[buf_idx, sample_idx, :local_seq_t, :local_vocab_t] + src = _rebuild_teacher_ipc_row(handle, device) dest[dest_seq, dest_vocab] = src[src_seq, src_vocab].to(torch.float32) return dest @@ -728,23 +856,80 @@ def _try_zero_copy_teacher_logits( student_seq_start = student_cp_rank * full_seq_len // student_cp_size student_seq_end = (student_cp_rank + 1) * full_seq_len // student_cp_size - # Exactly one full-vocab shard must cover this student rank's seq range. + from nemo_rl.models.policy.utils import validate_compact_teacher_ipc_handle + + # Exactly one physically stored full-vocab shard must cover this student + # rank's seq range. A compact zero suffix is logical coverage, but cannot be + # represented by a view and therefore takes the zero-filling fallback. chosen: list[dict[str, Any]] = [] for entry in per_sample_entries: - covering = [ - h - for h in entry["teacher_shards"] - if int(h["vocab_start_index"]) == 0 - and int(h["vocab_end_index"]) == full_vocab_size - and int(h["global_seq_start"]) <= student_seq_start - and int(h["global_seq_start"]) + int(h["actual_shape"][0]) - >= student_seq_end - ] + covering = [] + for handle in entry["teacher_shards"]: + compact_geometry = validate_compact_teacher_ipc_handle(handle) + stored_seq_len = ( + compact_geometry[1] + if compact_geometry is not None + else int(handle["actual_shape"][0]) + ) + if ( + int(handle["vocab_start_index"]) == 0 + and int(handle["vocab_end_index"]) == full_vocab_size + and int(handle["global_seq_start"]) <= student_seq_start + and int(handle["global_seq_start"]) + stored_seq_len >= student_seq_end + ): + covering.append(handle) if len(covering) != 1: return None chosen.append(covering[0]) - # All samples must form a contiguous slab in one storage slot. + compact_geometries = [ + validate_compact_teacher_ipc_handle(handle) for handle in chosen + ] + if any(geometry is not None for geometry in compact_geometries): + if any(geometry is None for geometry in compact_geometries): + raise ValueError( + "Dense teacher IPC zero-copy candidates mix compact and " + "rectangular storage layouts." + ) + h0 = chosen[0] + payload = h0["payload_ipc"] + teacher_seq_start = int(h0["global_seq_start"]) + seq_lo = student_seq_start - teacher_seq_start + seq_hi = student_seq_end - teacher_seq_start + first_geometry = compact_geometries[0] + assert first_geometry is not None + first_offset, first_stored_len, _, local_vocab_size = first_geometry + for index, (handle, geometry) in enumerate( + zip(chosen, compact_geometries, strict=True) + ): + assert geometry is not None + token_offset, stored_seq_len, _, handle_vocab_size = geometry + if ( + handle["payload_ipc"] != payload + or int(handle["global_seq_start"]) != teacher_seq_start + or tuple(handle["storage_shape"]) != tuple(h0["storage_shape"]) + or handle["dtype"] != h0["dtype"] + or handle_vocab_size != local_vocab_size + or stored_seq_len != first_stored_len + or token_offset != first_offset + index * first_stored_len + ): + return None + src_full, _ = _rebuild_compact_teacher_ipc_storage(h0, device) + if len(chosen) == 1: + return src_full[ + first_offset + seq_lo : first_offset + seq_hi, + :local_vocab_size, + ].unsqueeze(0) + return src_full.as_strided( + size=(len(chosen), seq_hi - seq_lo, local_vocab_size), + stride=(first_stored_len * local_vocab_size, local_vocab_size, 1), + storage_offset=( + src_full.storage_offset() + (first_offset + seq_lo) * local_vocab_size + ), + ) + + # Legacy rectangular samples must form a contiguous slab in one storage + # slot. h0 = chosen[0] payload = h0["payload_ipc"] buf_idx = int(h0["buf_idx"]) @@ -768,8 +953,14 @@ def rebuild_teacher_full_logits_from_ipc( per_sample_entries: list[dict[str, Any]], cp_group: Optional[torch.distributed.ProcessGroup], device: int, -) -> torch.Tensor: - """Rebuild ``[B, T_t/CP_s, V_t]`` teacher logits for this student rank. +) -> tuple[torch.Tensor, int]: + """Rebuild teacher logits and report actual fallback reconstructions. + + Returns ``([B, T_t/CP_s, V_t], fallback_count)``. ``fallback_count`` is + zero when the zero-copy view is used and otherwise counts the logical rows + that went through shard assembly. Reporting the value from this branch, + instead of predicting it from controller-side topology, keeps IPC + observability tied to the path that actually ran. Fast path (zero-copy view via :func:`_try_zero_copy_teacher_logits`): when the teacher is not vocab-sharded and each sample's seq range is covered by a @@ -793,7 +984,7 @@ def rebuild_teacher_full_logits_from_ipc( device=device, ) if view is not None: - return view + return view, 0 rebuilt = [ assemble_teacher_logits_from_shards( @@ -804,7 +995,1068 @@ def rebuild_teacher_full_logits_from_ipc( ) for entry in per_sample_entries ] - return torch.stack(rebuilt, dim=0) + # Packed CSD executes the loss at logical MBS1. Preserve the assembled + # row's storage in that overwhelmingly common path: torch.stack would + # allocate and copy another full [T_t/CP_s, V_t] tensor (about 1.24 GiB for + # Qwen3-14B at TP2/CP2) only to add a unit batch dimension. + if len(rebuilt) == 1: + return rebuilt[0].unsqueeze(0), 1 + return torch.stack(rebuilt, dim=0), len(rebuilt) + + +def rebuild_teacher_sparse_logits_from_ipc( + per_sample_entries: list[dict[str, Any]], + *, + device: int, +) -> SparseTeacherLogits: + """Rebuild full-sequence sparse teacher payloads from per-CP-shard IPC.""" + from nemo_rl.models.policy.utils import rebuild_cuda_tensor_from_ipc + + if not per_sample_entries: + raise ValueError("Sparse teacher IPC payload is empty.") + + per_sample_logits: list[torch.Tensor] = [] + per_sample_indices: list[torch.Tensor] = [] + per_sample_log_z: list[torch.Tensor] = [] + per_sample_gt_in_topk: list[torch.Tensor] = [] + has_gt_in_topk: Optional[bool] = None + + for sample_entry in per_sample_entries: + shards = sample_entry.get("teacher_shards", [sample_entry]) + shards = sorted(shards, key=lambda shard: int(shard["global_seq_start"])) + if not shards: + raise ValueError("Sparse teacher IPC sample has no sequence shards.") + shard_has_gt = ["gt_in_topk_ipc" in shard for shard in shards] + if len(set(shard_has_gt)) != 1: + raise ValueError( + "Sparse teacher IPC handles mix gt_in_topk and non-gt_in_topk " + "payloads within one sample." + ) + if has_gt_in_topk is None: + has_gt_in_topk = shard_has_gt[0] + elif has_gt_in_topk != shard_has_gt[0]: + raise ValueError( + "Sparse teacher IPC handles mix gt_in_topk and non-gt_in_topk " + "payloads across samples." + ) + + expected_start = 0 + for shard in shards: + shard_start = int(shard["global_seq_start"]) + if shard_start != expected_start: + raise ValueError( + "Sparse teacher CP shards must cover the sequence " + f"contiguously; expected start {expected_start}, got " + f"{shard_start}." + ) + expected_start += int(shard["topk_shape"][0]) + full_seq_len = int(shards[0]["full_seq_len"]) + if expected_start != full_seq_len: + raise ValueError( + "Sparse teacher CP shards do not cover the full sequence: " + f"covered={expected_start}, full_seq_len={full_seq_len}." + ) + + per_sample_logits.append( + torch.cat( + [ + rebuild_cuda_tensor_from_ipc( + shard["topk_logits_ipc"], device + ).detach() + for shard in shards + ], + dim=0, + ) + ) + per_sample_indices.append( + torch.cat( + [ + rebuild_cuda_tensor_from_ipc( + shard["topk_indices_ipc"], device + ).detach() + for shard in shards + ], + dim=0, + ) + ) + per_sample_log_z.append( + torch.cat( + [ + rebuild_cuda_tensor_from_ipc(shard["log_z_ipc"], device).detach() + for shard in shards + ], + dim=0, + ) + ) + if has_gt_in_topk: + per_sample_gt_in_topk.append( + torch.cat( + [ + rebuild_cuda_tensor_from_ipc( + shard["gt_in_topk_ipc"], device + ).detach() + for shard in shards + ], + dim=0, + ) + ) + + gt_in_topk = ( + torch.stack(per_sample_gt_in_topk, dim=0).to(torch.bool) + if has_gt_in_topk + else None + ) + return ( + torch.stack(per_sample_logits, dim=0).float(), + torch.stack(per_sample_indices, dim=0).to(torch.int32), + torch.stack(per_sample_log_z, dim=0).float(), + gt_in_topk, + ) + + +@dataclass(frozen=True) +class _SparseTeacherRoute: + """One global-sequence interval backed by one teacher CP shard.""" + + global_start: int + global_end: int + local_start: int + shard: dict[str, Any] + + +def _normalize_reader_device( + device: torch.device | int, +) -> tuple[torch.device, Optional[int]]: + """Return ``(device, cuda_index)``; the index is None for a CPU consumer.""" + if isinstance(device, int): + if device < 0: + raise ValueError(f"CUDA device id must be non-negative, got {device}") + return torch.device("cuda", device), device + if not isinstance(device, torch.device): + raise TypeError( + "device must be a torch.device or CUDA device id, got " + f"{type(device).__name__}" + ) + if device.type != "cuda": + return device, None + device_id = ( + torch.cuda.current_device() if device.index is None else int(device.index) + ) + return torch.device("cuda", device_id), device_id + + +def _reader_backing_key(field: str, value: Any) -> Any: + """Cache key for one IPC (or tensor) backing shared by several records.""" + if isinstance(value, torch.Tensor): + return (field, "tensor", id(value)) + # CUDA reduction handles are immutable tuples shared by every sample of a + # worker slab. Falling back to identity keeps unhashable test handles usable. + try: + hash(value) + except TypeError: + return (field, "ipc-id", id(value)) + return (field, "ipc", value) + + +def _open_reader_backing( + value: Any, *, device: torch.device, device_id: Optional[int] +) -> torch.Tensor: + """Map one IPC handle (or move one tensor) onto the consumer device.""" + if isinstance(value, torch.Tensor): + return value.detach().to(device) + if device_id is None: + raise ValueError( + f"CUDA IPC teacher backings require a CUDA consumer device, got {device}." + ) + from nemo_rl.models.policy.utils import rebuild_cuda_tensor_from_ipc + + return rebuild_cuda_tensor_from_ipc(value, device_id).detach() + + +def _sample_teacher_shards( + entry: Mapping[str, Any], sample_idx: int +) -> list[dict[str, Any]]: + shards = entry.get("teacher_shards", [entry]) + if not isinstance(shards, (list, tuple)): + raise TypeError( + "Teacher teacher_shards must be a list or tuple; " + f"sample={sample_idx}, got {type(shards).__name__}." + ) + if not shards: + raise ValueError(f"Teacher sample {sample_idx} contains no IPC shards") + if not all(isinstance(shard, dict) for shard in shards): + raise TypeError(f"Teacher sample {sample_idx} contains a non-dict shard") + return list(shards) + + +def _required_shard_int( + shard: dict[str, Any], field: str, *, sample_idx: int, shard_idx: int +) -> int: + try: + return int(shard[field]) + except KeyError as exc: + raise ValueError( + f"Sparse teacher shard is missing {field!r}; " + f"sample={sample_idx}, shard={shard_idx}." + ) from exc + except (TypeError, ValueError) as exc: + raise ValueError( + f"Sparse teacher {field!r} must be an integer; " + f"sample={sample_idx}, shard={shard_idx}." + ) from exc + + +class SparseTeacherRowReader: + """Cached router for sparse teacher rows kept in their teacher CP shards. + + Construct one reader per sparse teacher payload and reuse it for every + bounded row batch in a loss call. Construction validates and indexes the + compact ``sequence_segments`` metadata once. CUDA IPC tensors are opened + lazily and cached, so repeated row batches neither reopen the same backing + storage nor scan all CP shards for every row. :meth:`close` drops the + cached backings. + + Each shard record carries natural top-k logits and ids sorted by id, the + exact log Z, the score-ranked k-th id, and a forced-label sidecar (see + :meth:`MegatronPolicyWorker.get_sparse_logits_ipc`). Tensor-backed records + use the same path and may target CPU. + """ + + TRANSPORT = "sparse_topk" + _BACKING_FIELDS = ( + "topk_logits_ipc", + "topk_indices_ipc", + "log_z_ipc", + "natural_tail_indices_ipc", + "forced_logits_ipc", + "forced_indices_ipc", + "forced_in_topk_ipc", + ) + + def __init__( + self, + per_sample_entries: Sequence[dict[str, Any]], + *, + device: torch.device | int, + ) -> None: + if not per_sample_entries: + raise ValueError("Sparse teacher payload contains no samples") + consumer_device, consumer_device_id = _normalize_reader_device(device) + self.device: torch.device = consumer_device + self._device_id: Optional[int] = consumer_device_id + self._backing_tensors: dict[Any, torch.Tensor] = {} + self._routes: list[list[_SparseTeacherRoute]] = [] + self._route_starts: list[list[int]] = [] + self.k: int + self.temperature: float + self.real_vocab_size: int + self.full_seq_len: int + self.force_width: int + self.membership_k: int + self._build_routing_index(per_sample_entries) + + def _build_routing_index( + self, per_sample_entries: Sequence[dict[str, Any]] + ) -> None: + expected_meta: Optional[tuple[int, float, int, int]] = None + expected_full_seq_len: Optional[int] = None + expected_force_width: Optional[int] = None + + for sample_idx, entry in enumerate(per_sample_entries): + if not isinstance(entry, dict): + raise TypeError( + "Sparse teacher samples must be dictionaries; " + f"sample={sample_idx}, got {type(entry).__name__}." + ) + sample_routes: list[_SparseTeacherRoute] = [] + declared_full_seq_len: Optional[int] = None + + for shard_idx, shard in enumerate( + _sample_teacher_shards(entry, sample_idx) + ): + transport = shard.get("transport") + if transport != self.TRANSPORT: + raise ValueError( + "Unsupported sparse teacher transport " + f"{transport!r}; sample={sample_idx}, shard={shard_idx}." + ) + k = _required_shard_int( + shard, "k", sample_idx=sample_idx, shard_idx=shard_idx + ) + real_vocab_size = _required_shard_int( + shard, "real_vocab_size", sample_idx=sample_idx, shard_idx=shard_idx + ) + try: + temperature = float(shard["temperature"]) + except KeyError as exc: + raise ValueError( + "Sparse teacher shard is missing 'temperature'; " + f"sample={sample_idx}, shard={shard_idx}." + ) from exc + except (TypeError, ValueError) as exc: + raise ValueError( + "Sparse teacher temperature must be numeric; " + f"sample={sample_idx}, shard={shard_idx}." + ) from exc + if k <= 0 or k > real_vocab_size: + raise ValueError( + "Sparse teacher k must satisfy 0 < k <= real_vocab_size; " + f"got k={k}, real_vocab_size={real_vocab_size}." + ) + if not math.isfinite(temperature) or temperature <= 0: + raise ValueError( + "Sparse teacher temperature must be finite and > 0, got " + f"{temperature}." + ) + membership_k = _required_shard_int( + shard, + "forced_in_topk_k", + sample_idx=sample_idx, + shard_idx=shard_idx, + ) + force_width = _required_shard_int( + shard, "force_width", sample_idx=sample_idx, shard_idx=shard_idx + ) + if force_width <= 0: + raise ValueError( + f"Sparse teacher force_width must be positive; got {force_width}." + ) + if expected_force_width is None: + expected_force_width = force_width + elif force_width != expected_force_width: + raise ValueError( + "Sparse teacher CP shards disagree on force_width." + ) + if membership_k <= 0 or membership_k > real_vocab_size: + raise ValueError( + "Sparse teacher membership top-k must satisfy " + "0 < membership_k <= real_vocab_size; " + f"got {membership_k}." + ) + metadata = (k, temperature, real_vocab_size, membership_k) + if expected_meta is None: + expected_meta = metadata + elif metadata != expected_meta: + raise ValueError( + "Sparse teacher CP shards disagree on " + "k/temperature/vocabulary/membership support." + ) + + missing_backings = [ + field for field in self._BACKING_FIELDS if field not in shard + ] + if missing_backings: + raise ValueError( + "Sparse teacher shard is missing backing fields " + f"{missing_backings}; sample={sample_idx}, shard={shard_idx}." + ) + buf_idx = _required_shard_int( + shard, "buf_idx", sample_idx=sample_idx, shard_idx=shard_idx + ) + sample_index_in_buf = _required_shard_int( + shard, + "sample_index_in_buf", + sample_idx=sample_idx, + shard_idx=shard_idx, + ) + if buf_idx < 0 or sample_index_in_buf < 0: + raise ValueError( + "Sparse teacher buf_idx/sample_index_in_buf must be " + f"non-negative; sample={sample_idx}, shard={shard_idx}." + ) + + shape = shard.get("topk_shape") + if not isinstance(shape, (list, tuple)) or len(shape) != 2: + raise ValueError( + "Sparse teacher topk_shape must be a two-element sequence; " + f"sample={sample_idx}, shard={shard_idx}." + ) + try: + local_seq_len, shard_k = (int(x) for x in shape) + except (TypeError, ValueError) as exc: + raise ValueError( + "Sparse teacher topk_shape entries must be integers; " + f"sample={sample_idx}, shard={shard_idx}." + ) from exc + if local_seq_len <= 0 or shard_k != k: + raise ValueError( + "Sparse teacher topk_shape disagrees with shard metadata; " + f"shape={(local_seq_len, shard_k)}, k={k}." + ) + forced_shape = shard.get("forced_shape") + if ( + not isinstance(forced_shape, (list, tuple)) + or len(forced_shape) != 2 + ): + raise ValueError( + "Sparse teacher forced_shape must be a two-element sequence; " + f"sample={sample_idx}, shard={shard_idx}." + ) + try: + forced_seq_len, shape_force_width = (int(x) for x in forced_shape) + except (TypeError, ValueError) as exc: + raise ValueError( + "Sparse teacher forced_shape entries must be integers; " + f"sample={sample_idx}, shard={shard_idx}." + ) from exc + if forced_seq_len != local_seq_len or shape_force_width != force_width: + raise ValueError( + "Sparse teacher forced_shape disagrees with topk " + "sequence/force_width metadata; " + f"forced_shape={(forced_seq_len, shape_force_width)}, " + f"topk_seq={local_seq_len}, force_width={force_width}." + ) + scalar_shape = shard.get("scalar_shape") + if scalar_shape is not None and tuple(scalar_shape) != (local_seq_len,): + raise ValueError( + "Sparse teacher scalar_shape must match topk sequence " + f"length {local_seq_len}, got {scalar_shape}." + ) + + shard_full_seq = shard.get("full_seq_len") + if shard_full_seq is not None: + shard_full_seq = int(shard_full_seq) + if shard_full_seq <= 0: + raise ValueError( + "Sparse teacher full_seq_len must be positive, got " + f"{shard_full_seq}." + ) + if declared_full_seq_len is None: + declared_full_seq_len = shard_full_seq + elif shard_full_seq != declared_full_seq_len: + raise ValueError( + "Sparse teacher shards in one sample disagree on " + "full_seq_len." + ) + + raw_segments = shard.get("sequence_segments") + if not isinstance(raw_segments, (list, tuple)) or not raw_segments: + raise ValueError( + "Sparse teacher shard must contain non-empty " + f"sequence_segments; sample={sample_idx}, shard={shard_idx}." + ) + local_intervals: list[tuple[int, int]] = [] + for segment_idx, segment in enumerate(raw_segments): + if not isinstance(segment, (list, tuple)) or len(segment) != 3: + raise ValueError( + "Sparse teacher sequence segments must be " + "(local_start, global_start, length) triples; " + f"sample={sample_idx}, shard={shard_idx}, " + f"segment={segment_idx}." + ) + try: + local_start, global_start, length = (int(x) for x in segment) + except (TypeError, ValueError) as exc: + raise ValueError( + "Sparse teacher sequence segment entries must be " + f"integers; sample={sample_idx}, shard={shard_idx}." + ) from exc + if local_start < 0 or global_start < 0 or length <= 0: + raise ValueError( + "Sparse teacher sequence segments require non-negative " + "starts and positive lengths; " + f"got {(local_start, global_start, length)}." + ) + local_end = local_start + length + global_end = global_start + length + if local_end > local_seq_len: + raise ValueError( + "Sparse teacher sequence segment exceeds topk_shape; " + f"local_end={local_end}, sequence={local_seq_len}." + ) + if ( + declared_full_seq_len is not None + and global_end > declared_full_seq_len + ): + raise ValueError( + "Sparse teacher sequence segment exceeds full_seq_len; " + f"global_end={global_end}, full={declared_full_seq_len}." + ) + local_intervals.append((local_start, local_end)) + sample_routes.append( + _SparseTeacherRoute( + global_start=global_start, + global_end=global_end, + local_start=local_start, + shard=shard, + ) + ) + + local_intervals.sort() + local_cursor = 0 + for local_start, local_end in local_intervals: + if local_start != local_cursor: + raise ValueError( + "Sparse teacher shard sequence_segments must cover " + "each local row exactly once; " + f"expected local start {local_cursor}, got {local_start}." + ) + local_cursor = local_end + if local_cursor != local_seq_len: + raise ValueError( + "Sparse teacher shard sequence_segments do not cover the " + f"local sequence: covered={local_cursor}, " + f"sequence={local_seq_len}." + ) + + sample_routes.sort(key=lambda route: route.global_start) + global_cursor = 0 + for route in sample_routes: + if route.global_start != global_cursor: + raise ValueError( + "Sparse teacher CP coverage must cover each global " + "position exactly once; " + f"sample={sample_idx}, expected start {global_cursor}, " + f"got {route.global_start}." + ) + global_cursor = route.global_end + sample_full_seq_len = declared_full_seq_len or global_cursor + if global_cursor != sample_full_seq_len: + raise ValueError( + "Sparse teacher CP shards do not cover the full sequence; " + f"sample={sample_idx}, covered={global_cursor}, " + f"full_seq_len={sample_full_seq_len}." + ) + if expected_full_seq_len is None: + expected_full_seq_len = sample_full_seq_len + elif sample_full_seq_len != expected_full_seq_len: + raise ValueError( + "Sparse teacher samples disagree on full sequence length; " + f"expected {expected_full_seq_len}, got {sample_full_seq_len}." + ) + self._routes.append(sample_routes) + self._route_starts.append([route.global_start for route in sample_routes]) + + assert ( + expected_meta is not None + and expected_full_seq_len is not None + and expected_force_width is not None + ) + self.k, self.temperature, self.real_vocab_size, self.membership_k = ( + expected_meta + ) + self.force_width = expected_force_width + self.full_seq_len = expected_full_seq_len + + def _backing(self, shard: dict[str, Any], field: str) -> torch.Tensor: + value = shard[field] + key = _reader_backing_key(field, value) + tensor = self._backing_tensors.get(key) + if tensor is None: + tensor = _open_reader_backing( + value, device=self.device, device_id=self._device_id + ) + self._backing_tensors[key] = tensor + + seq_len, width = (int(x) for x in shard["topk_shape"]) + if field in ("topk_logits_ipc", "topk_indices_ipc"): + expected_width = width + elif field in ("forced_logits_ipc", "forced_indices_ipc", "forced_in_topk_ipc"): + expected_width = int(shard["force_width"]) + else: + expected_width = 1 + buf_idx = int(shard["buf_idx"]) + sample_idx = int(shard["sample_index_in_buf"]) + if tensor.dim() != 4 or ( + tensor.shape[0] <= buf_idx + or tensor.shape[1] <= sample_idx + or tensor.shape[2] < seq_len + or tensor.shape[3] < expected_width + ): + raise ValueError( + f"Sparse teacher {field} backing shape {tuple(tensor.shape)} " + "does not cover declared slot/shape " + f"({buf_idx}, {sample_idx}, {seq_len}, {expected_width})." + ) + if field in ( + "topk_indices_ipc", + "natural_tail_indices_ipc", + "forced_indices_ipc", + ): + if tensor.is_floating_point() or tensor.dtype == torch.bool: + raise TypeError( + f"Sparse teacher {field} backing must have an integer " + f"dtype, got {tensor.dtype}." + ) + elif field == "forced_in_topk_ipc": + if tensor.dtype != torch.bool: + raise TypeError( + f"Sparse teacher {field} backing must have bool dtype, " + f"got {tensor.dtype}." + ) + elif not tensor.is_floating_point(): + raise TypeError( + f"Sparse teacher {field} backing must be floating point, " + f"got {tensor.dtype}." + ) + return tensor + + def _route(self, batch_idx: int, global_position: int) -> _SparseTeacherRoute: + if not 0 <= batch_idx < len(self._routes): + raise IndexError( + f"Sparse teacher batch index {batch_idx} outside " + f"[0, {len(self._routes)})" + ) + route_idx = bisect_right(self._route_starts[batch_idx], global_position) - 1 + if ( + not 0 <= global_position < self.full_seq_len + or route_idx < 0 + or global_position >= self._routes[batch_idx][route_idx].global_end + ): + raise ValueError( + "Sparse teacher CP coverage must resolve each requested position " + f"exactly once; batch={batch_idx}, pos={global_position}, matches=0." + ) + return self._routes[batch_idx][route_idx] + + def gather_rows( + self, + batch_indices: torch.Tensor, + global_positions: torch.Tensor, + required_token_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Gather one exact K-wide support per loss request. + + Every request names the teacher label it needs. When that label is + outside the natural top-k, the producer-recorded score-rank-k token is + evicted and only the requested sidecar label is inserted, so two + requests for the same row may receive different supports without + changing the backing row. Requests are grouped by backing shard, so + each source is indexed once per call. + + Args: + batch_indices: Batch rows ``[N]``. + global_positions: Global teacher sequence positions ``[N]``. + required_token_ids: Teacher label per request ``[N]``; must be one + of the forced labels exported for that row. + + Returns: + ``(topk_logits [N, k] fp32, topk_indices [N, k] int32 sorted by id, + log_z [N] fp32, requested_in_topk [N] bool)``. ``topk_logits`` are + raw (not temperature scaled); ``log_z`` is the temperature-scaled + full-vocabulary normalizer; ``requested_in_topk`` is the label's + membership in the producer's natural top-``membership_k``. + """ + if ( + batch_indices.dim() != 1 + or global_positions.dim() != 1 + or required_token_ids.dim() != 1 + ): + raise ValueError( + "batch_indices, global_positions and required_token_ids must be " + "one-dimensional" + ) + if not ( + batch_indices.shape == global_positions.shape == required_token_ids.shape + ): + raise ValueError( + "batch_indices, global_positions and required_token_ids must have " + f"matching shapes; got {tuple(batch_indices.shape)}, " + f"{tuple(global_positions.shape)} and {tuple(required_token_ids.shape)}" + ) + for name, tensor in ( + ("batch_indices", batch_indices), + ("global_positions", global_positions), + ("required_token_ids", required_token_ids), + ): + if tensor.dtype == torch.bool or tensor.is_floating_point(): + raise TypeError(f"{name} must have an integer dtype") + requests = list( + zip( + batch_indices.detach().cpu().tolist(), + global_positions.detach().cpu().tolist(), + ) + ) + if not requests: + raise ValueError("At least one sparse teacher row must be requested") + + n = len(requests) + topk_logits = torch.empty((n, self.k), device=self.device, dtype=torch.float32) + topk_indices = torch.empty((n, self.k), device=self.device, dtype=torch.int32) + log_z = torch.empty((n,), device=self.device, dtype=torch.float32) + natural_tail_indices = torch.empty((n,), device=self.device, dtype=torch.int32) + forced_logits = torch.empty( + (n, self.force_width), device=self.device, dtype=torch.float32 + ) + forced_indices = torch.empty( + (n, self.force_width), device=self.device, dtype=torch.int32 + ) + forced_in_topk = torch.empty( + (n, self.force_width), device=self.device, dtype=torch.bool + ) + required = required_token_ids.detach().to(device=self.device, dtype=torch.long) + invalid_required = (required < 0) | (required >= self.real_vocab_size) + if bool(invalid_required.any().item()): + bad_id = int(required[invalid_required][0].item()) + raise ValueError( + "required_token_ids must contain real-vocabulary ids in " + f"[0, {self.real_vocab_size}); got {bad_id}." + ) + + grouped: dict[int, tuple[dict[str, Any], list[int], list[int]]] = {} + for out_idx, (batch_idx, global_position) in enumerate(requests): + route = self._route(int(batch_idx), int(global_position)) + # Each sample/shard record identifies one backing slot; group its + # load-balanced segments together without hashing the record. + group = grouped.get(id(route.shard)) + if group is None: + group = (route.shard, [], []) + grouped[id(route.shard)] = group + group[1].append(out_idx) + group[2].append( + route.local_start + int(global_position) - route.global_start + ) + + destinations = ( + ("topk_logits_ipc", topk_logits, self.k), + ("topk_indices_ipc", topk_indices, self.k), + ("log_z_ipc", log_z, None), + ("natural_tail_indices_ipc", natural_tail_indices, None), + ("forced_logits_ipc", forced_logits, self.force_width), + ("forced_indices_ipc", forced_indices, self.force_width), + ("forced_in_topk_ipc", forced_in_topk, self.force_width), + ) + for shard, output_rows, local_rows in grouped.values(): + buf_idx = int(shard["buf_idx"]) + sample_idx = int(shard["sample_index_in_buf"]) + output_index = torch.tensor( + output_rows, device=self.device, dtype=torch.long + ) + local_index = torch.tensor(local_rows, device=self.device, dtype=torch.long) + for field, destination, width in destinations: + source = self._backing(shard, field)[buf_idx, sample_idx].index_select( + 0, local_index + ) + source = source[:, :width] if width is not None else source[:, 0] + destination.index_copy_(0, output_index, source.to(destination.dtype)) + + if bool(((topk_indices < 0) | (topk_indices >= self.real_vocab_size)).any()): + raise ValueError("Sparse teacher natural top-k contains an invalid id") + if self.k > 1 and bool((topk_indices[:, 1:] <= topk_indices[:, :-1]).any()): + raise ValueError("Sparse teacher natural top-k ids must be strictly sorted") + if bool( + ((forced_indices < -1) | (forced_indices >= self.real_vocab_size)).any() + ): + raise ValueError("Sparse teacher forced sidecar contains an invalid id") + for left in range(self.force_width): + for right in range(left + 1, self.force_width): + duplicate = (forced_indices[:, left] >= 0) & ( + forced_indices[:, left] == forced_indices[:, right] + ) + if bool(duplicate.any().item()): + raise ValueError( + "Sparse teacher forced sidecar contains duplicate labels" + ) + + sidecar_match = forced_indices.long() == required.unsqueeze(-1) + if bool((sidecar_match.sum(dim=-1) != 1).any().item()): + raise ValueError( + "Sparse teacher required label must appear exactly once in the " + "forced sidecar" + ) + selected_forced_logits = ( + torch.where( + sidecar_match, forced_logits, torch.full_like(forced_logits, -torch.inf) + ) + .max(dim=-1) + .values + ) + requested_in_topk = (sidecar_match & forced_in_topk).any(dim=-1) + + natural_match = topk_indices.long() == required.unsqueeze(-1) + if bool((natural_match.sum(dim=-1) > 1).any().item()): + raise ValueError("Sparse teacher natural top-k contains duplicate ids") + natural_requested_logits = torch.where( + natural_match, topk_logits, torch.zeros_like(topk_logits) + ).sum(dim=-1) + inconsistent_duplicate = natural_match.any(dim=-1) & ( + natural_requested_logits != selected_forced_logits + ) + if bool(inconsistent_duplicate.any().item()): + raise ValueError( + "Sparse teacher natural and forced copies of a requested label " + "have different raw logits" + ) + needs_insert = ~natural_match.any(dim=-1) + + tail_match = topk_indices.long() == natural_tail_indices.long().unsqueeze(-1) + if bool((tail_match.sum(dim=-1) != 1).any().item()): + raise ValueError( + "Sparse teacher natural tail id must appear exactly once in top-k" + ) + evict_slot = tail_match.to(torch.long).argmax(dim=-1) + materialized_logits = topk_logits.clone() + materialized_indices = topk_indices.clone() + insert_rows = needs_insert.nonzero(as_tuple=False).flatten() + if insert_rows.numel() > 0: + insert_slots = evict_slot.index_select(0, insert_rows) + materialized_logits[insert_rows, insert_slots] = ( + selected_forced_logits.index_select(0, insert_rows) + ) + materialized_indices[insert_rows, insert_slots] = required.index_select( + 0, insert_rows + ).to(materialized_indices.dtype) + + sorted_indices, order = torch.sort(materialized_indices, dim=-1) + sorted_logits = torch.gather(materialized_logits, dim=-1, index=order) + if self.k > 1 and bool((sorted_indices[:, 1:] <= sorted_indices[:, :-1]).any()): + raise ValueError( + "Sparse teacher materialized top-k ids must be strictly sorted" + ) + return sorted_logits, sorted_indices, log_z, requested_in_topk + + def close(self) -> None: + """Drop the cached backing tensors (and their imported IPC storage).""" + self._backing_tensors.clear() + + +def gather_sparse_teacher_rows_from_ipc( + per_sample_entries: Sequence[dict[str, Any]], + batch_indices: torch.Tensor, + global_positions: torch.Tensor, + required_token_ids: torch.Tensor, + *, + device: torch.device | int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Read only the requested rows from sparse teacher CP shards. + + One-shot form of :meth:`SparseTeacherRowReader.gather_rows`. Each + ``(batch, global_position)`` resolves to exactly one ``sequence_segments`` + interval; only those rows are copied, and the teacher CP shards are never + concatenated into a full sequence. + """ + reader = SparseTeacherRowReader(per_sample_entries, device=device) + try: + rows = reader.gather_rows(batch_indices, global_positions, required_token_ids) + finally: + reader.close() + return rows + + +def _teacher_ipc_row_view( + handle: Mapping[str, Any], src_full: torch.Tensor +) -> torch.Tensor: + """Return the physically stored ``[T_stored, V_local]`` row of ``handle``. + + ``src_full`` is the handle's opened producer storage: a compact + ``[tokens, V_local]`` slab or a rectangular ``[N_mb, B, T, V_local]`` one. + """ + from nemo_rl.models.policy.utils import validate_compact_teacher_ipc_handle + + compact_geometry = validate_compact_teacher_ipc_handle(handle) + if compact_geometry is not None: + token_offset, stored_seq_len, _used_tokens, local_vocab_size = compact_geometry + if stored_seq_len <= 0: + raise ValueError("A zero-length compact IPC row has no physical view.") + expected_storage_shape = tuple(int(value) for value in handle["storage_shape"]) + if tuple(src_full.shape) != expected_storage_shape or src_full.ndim != 2: + raise ValueError( + "Compact dense teacher IPC rebuilt an unexpected storage shape: " + f"actual={tuple(src_full.shape)}, expected={expected_storage_shape}." + ) + if not src_full.is_contiguous(): + raise ValueError("Compact dense teacher IPC storage must be contiguous.") + if src_full.dtype != handle["dtype"]: + raise ValueError( + "Compact dense teacher IPC rebuilt an unexpected dtype: " + f"actual={src_full.dtype}, expected={handle['dtype']}." + ) + return src_full[token_offset : token_offset + stored_seq_len, :local_vocab_size] + + actual_shape = handle["actual_shape"] + if ( + not isinstance(actual_shape, (list, tuple, torch.Size)) + or len(actual_shape) != 2 + ): + raise ValueError( + "Dense teacher IPC actual_shape must be two-dimensional, got " + f"{actual_shape!r}." + ) + local_seq_len, local_vocab_size = (int(value) for value in actual_shape) + buf_idx = int(handle["buf_idx"]) + sample_idx = int(handle["sample_index_in_buf"]) + if ( + src_full.ndim != 4 + or buf_idx < 0 + or buf_idx >= src_full.shape[0] + or sample_idx < 0 + or sample_idx >= src_full.shape[1] + or local_seq_len <= 0 + or local_seq_len > src_full.shape[2] + or local_vocab_size <= 0 + or local_vocab_size > src_full.shape[3] + ): + raise ValueError( + "Dense teacher IPC rectangular handle is outside rebuilt storage: " + f"storage={tuple(src_full.shape)}, buf_idx={buf_idx}, " + f"sample_idx={sample_idx}, actual_shape={tuple(actual_shape)}." + ) + return src_full[buf_idx, sample_idx, :local_seq_len, :local_vocab_size] + + +@dataclass(frozen=True) +class _DenseTeacherRoute: + """One teacher CP sequence window and its TP vocabulary shards.""" + + seq_start: int + seq_end: int + shards: tuple[dict[str, Any], ...] + + +class DenseTeacherRowReader: + """Read requested full-vocabulary rows from dense teacher IPC shards. + + The dense counterpart of :class:`SparseTeacherRowReader`: instead of + assembling a teacher's CP-local ``[B, T/CP, V]`` window, a loss reads only + the rows it needs. Each requested row is stitched across the teacher's TP + vocabulary shards into ``[N, V]`` fp32. Rows inside a compact handle's + omitted (zero) suffix read as zeros, matching dense reassembly. Every + payload is opened once and cached until :meth:`close`. + """ + + def __init__( + self, + per_sample_entries: Sequence[dict[str, Any]], + *, + device: torch.device | int, + ) -> None: + from nemo_rl.models.policy.utils import _validate_dense_ipc_shard_coverage + + if not per_sample_entries: + raise ValueError("Dense teacher payload contains no samples") + consumer_device, consumer_device_id = _normalize_reader_device(device) + self.device: torch.device = consumer_device + self._device_id: Optional[int] = consumer_device_id + self._storage: dict[Any, torch.Tensor] = {} + self._routes: list[list[_DenseTeacherRoute]] = [] + self._route_starts: list[list[int]] = [] + full_seq_len: Optional[int] = None + full_vocab_size: Optional[int] = None + for sample_idx, entry in enumerate(per_sample_entries): + shards = _sample_teacher_shards(entry, sample_idx) + _validate_dense_ipc_shard_coverage(sample_idx, shards) + sample_seq_len = int(shards[0]["full_seq_len"]) + sample_vocab_size = int(shards[0]["full_vocab_size"]) + if full_seq_len is None: + full_seq_len, full_vocab_size = sample_seq_len, sample_vocab_size + elif (sample_seq_len, sample_vocab_size) != (full_seq_len, full_vocab_size): + raise ValueError( + "Dense teacher samples disagree on the logical [T, V] shape: " + f"expected {(full_seq_len, full_vocab_size)}, got " + f"{(sample_seq_len, sample_vocab_size)}." + ) + by_window: dict[tuple[int, int], list[dict[str, Any]]] = {} + for shard in shards: + seq_start = int(shard["global_seq_start"]) + seq_end = seq_start + int(shard["actual_shape"][0]) + by_window.setdefault((seq_start, seq_end), []).append(shard) + routes = [ + _DenseTeacherRoute( + seq_start=seq_start, + seq_end=seq_end, + shards=tuple( + sorted(window, key=lambda s: int(s["vocab_start_index"])) + ), + ) + for (seq_start, seq_end), window in sorted(by_window.items()) + ] + self._routes.append(routes) + self._route_starts.append([route.seq_start for route in routes]) + assert full_seq_len is not None and full_vocab_size is not None + self.full_seq_len = full_seq_len + self.full_vocab_size = full_vocab_size + + def _row_view(self, handle: dict[str, Any]) -> torch.Tensor: + payload = handle["payload_ipc"] + key = _reader_backing_key("payload_ipc", payload) + src_full = self._storage.get(key) + if src_full is None: + src_full = _open_reader_backing( + payload, device=self.device, device_id=self._device_id + ) + self._storage[key] = src_full + return _teacher_ipc_row_view(handle, src_full) + + def gather_rows( + self, batch_indices: torch.Tensor, global_positions: torch.Tensor + ) -> torch.Tensor: + """Return raw teacher logits ``[N, full_vocab_size]`` (fp32) for the requests. + + Args: + batch_indices: Batch rows ``[N]``. + global_positions: Global teacher sequence positions ``[N]``. + """ + from nemo_rl.models.policy.utils import validate_compact_teacher_ipc_handle + + if batch_indices.dim() != 1 or batch_indices.shape != global_positions.shape: + raise ValueError( + "batch_indices and global_positions must be matching " + f"one-dimensional tensors; got {tuple(batch_indices.shape)} and " + f"{tuple(global_positions.shape)}" + ) + requests = list( + zip( + batch_indices.detach().cpu().tolist(), + global_positions.detach().cpu().tolist(), + ) + ) + out = torch.zeros( + (len(requests), self.full_vocab_size), + device=self.device, + dtype=torch.float32, + ) + grouped: dict[tuple[int, int], tuple[_DenseTeacherRoute, list[int], list[int]]] + grouped = {} + for out_idx, (batch_idx, position) in enumerate(requests): + batch_idx, position = int(batch_idx), int(position) + if not 0 <= batch_idx < len(self._routes): + raise IndexError( + f"Dense teacher batch index {batch_idx} outside " + f"[0, {len(self._routes)})" + ) + if not 0 <= position < self.full_seq_len: + raise IndexError( + f"Dense teacher position {position} outside " + f"[0, {self.full_seq_len})" + ) + route_idx = bisect_right(self._route_starts[batch_idx], position) - 1 + route = self._routes[batch_idx][route_idx] + group = grouped.setdefault((batch_idx, route_idx), (route, [], [])) + group[1].append(out_idx) + group[2].append(position - route.seq_start) + + for route, output_rows, local_rows in grouped.values(): + for shard in route.shards: + compact_geometry = validate_compact_teacher_ipc_handle(shard) + stored_seq_len = ( + compact_geometry[1] + if compact_geometry is not None + else int(shard["actual_shape"][0]) + ) + keep = [ + (out_row, local_row) + for out_row, local_row in zip(output_rows, local_rows) + if local_row < stored_seq_len + ] + if not keep: + continue + stored = self._row_view(shard) + output_index = torch.tensor( + [row for row, _ in keep], device=self.device, dtype=torch.long + ) + local_index = torch.tensor( + [row for _, row in keep], device=stored.device, dtype=torch.long + ) + vocab_start = int(shard["vocab_start_index"]) + vocab_end = int(shard["vocab_end_index"]) + out[output_index, vocab_start:vocab_end] = ( + stored.index_select(0, local_index).to(out.device).float() + ) + return out + + def close(self) -> None: + """Drop the cached producer storages (and their imported IPC handles).""" + self._storage.clear() def valid_chunk_mask( @@ -1058,7 +2310,7 @@ def build_exact_token_map( student row's projection weights descending, then picks an exact-token map per the ``xtoken_loss`` flag: - - ``xtoken_loss=False`` (strict): ``has_exact_map = (sorted_values[:, 0] == 1.0) & (projection_indices[:, 1] == -1)``. + - ``xtoken_loss=False`` (strict): ``has_exact_map = sorted_values[:, 0] == 1.0``. On collision (multiple students mapping to the same teacher id), the earliest (lowest) student index wins. - ``xtoken_loss=True`` (relaxed): ``has_exact_map = sorted_values[:, 0] >= 0.6``. @@ -1096,11 +2348,18 @@ def build_exact_token_map( if xtoken_loss: has_exact_map = sorted_values[:, 0] >= 0.6 else: - # Strict: exactly one top-k entry with weight 1.0, no second - # mapping. `indices[:, 1] == -1` is the sentinel used by the - # `_exact_map_remapped` projection files for "no second - # mapping". - has_exact_map = (sorted_values[:, 0] == 1.0) & (indices[:, 1] == -1) + # Strict: top-k entry with weight 1.0 (no probability mass leaks to + # other teacher ids — the student token maps 1-to-1). The legacy + # `_exact_map_remapped` files used `indices[:, 1] == -1` as a sentinel + # for "no second mapping", but modern projection builders + # (tools/x_token/minimal_projection_via_multitoken.py with + # --enable-exact-match) leave indices[:, 1+] filled with valid ids whose + # weight is 0. Do NOT re-add the sentinel guard: it silently empties the + # common vocab for every such matrix (e.g. nemotron-nano-3.5 <- Qwen3-30B + # went 68131 -> 0 common tokens, routing all 496k 1-to-1 chunks down the + # mismatch path). The top1 == 1.0 check is sufficient, since by + # normalization weight 1.0 on one entry means 0 elsewhere. + has_exact_map = sorted_values[:, 0] == 1.0 # Gather (s_idx, t_idx, prob) for each exact-map candidate. s_candidates = torch.where(has_exact_map)[0] @@ -1172,6 +2431,562 @@ def build_exact_token_map( return result +def _chunk_ids_to_spans(chunk_id: torch.Tensor, max_pairs: int) -> torch.Tensor: + """Contiguous ``[start, end)`` position span per chunk id. + + Args: + chunk_id: ``[B, T]`` where each entry is the pair/chunk index the + position belongs to (sentinel ``< 0`` or ``>= max_pairs`` for + unaligned positions). + max_pairs: chunk-id space width (the ``pair_valid`` slot count). + + Returns: + ``[B, max_pairs, 2]`` long tensor; row ``k`` is ``[first, last + 1)`` + over the positions with ``chunk_id == k``. A chunk id absent from a + sample gets ``[0, 0)`` (zero-length — the CSD ``M/N > 0`` gate and + ``pair_valid`` both skip it). Assumes each chunk's positions are + contiguous (the offset-based cluster alignment guarantees this). + """ + b, t = chunk_id.shape + device = chunk_id.device + pos = torch.arange(t, device=device).unsqueeze(0).expand(b, t) + valid = (chunk_id >= 0) & (chunk_id < max_pairs) + idx = chunk_id.clamp(0, max_pairs - 1) + # Invalid positions carry T (ignored by amin) / -1 (ignored by amax), so a + # clamped-in bucket is never corrupted by an unaligned position. + pos_for_min = torch.where(valid, pos, torch.full_like(pos, t)) + pos_for_max = torch.where(valid, pos, torch.full_like(pos, -1)) + first = torch.full((b, max_pairs), t, dtype=torch.long, device=device) + last = torch.full((b, max_pairs), -1, dtype=torch.long, device=device) + first.scatter_reduce_(1, idx, pos_for_min, reduce="amin", include_self=True) + last.scatter_reduce_(1, idx, pos_for_max, reduce="amax", include_self=True) + has = first < t + spans = torch.zeros(b, max_pairs, 2, dtype=torch.long, device=device) + spans[:, :, 0] = torch.where(has, first, torch.zeros_like(first)) + spans[:, :, 1] = torch.where(has, last + 1, torch.zeros_like(last)) + return spans + + +def loss_replica_group( + cp_group: Optional[torch.distributed.ProcessGroup], +) -> Optional[torch.distributed.ProcessGroup]: + """Group over which to reduce the loss's global normalizers, or ``None``. + + The loss needs a "sum one contribution per distinct sample shard" group. The + obvious choice, ``torch.distributed.group.WORLD``, **deadlocks under pipeline + parallelism**: only the last stage runs the loss, so a WORLD collective is + never joined by the earlier stages and every rank hangs until the NCCL + watchdog fires. + + Megatron's data-parallel(-with-CP) group is the right group instead: its + members all share the same pipeline-stage / tensor-parallel coordinates, so + it is confined to the stage that actually runs the loss, and it spans exactly + the DP x CP axes the normalizers want. Returns ``None`` when Megatron is not + the active backend (DTensor, or a single-process / CPU test), leaving callers + on their existing WORLD path. + """ + if cp_group is None or not torch.distributed.is_initialized(): + return None + try: + # Local import keeps the optional Megatron dependency boundary intact; + # a non-Megatron caller falls back to None. + from megatron.core import parallel_state + except ImportError: + return None + if not parallel_state.is_initialized(): + return None + return parallel_state.get_data_parallel_group(with_context_parallel=True) + + +def _student_seq_to_contiguous_window( + x: torch.Tensor, + cp_group: Optional[torch.distributed.ProcessGroup], + *, + data_is_cp_sharded: bool, + seq_dim: int = 1, +) -> torch.Tensor: + """Relay a student-sequence data tensor to this CP rank's contiguous window. + + The two training backends hand the loss different layouts for the + student-sequence entries of the microbatch dict: + + * DTensor worker: the buffers were CP-sharded in place by PyTorch's + ``context_parallel``, i.e. they are in the *load-balanced* (``2*cp`` + interleaved) layout and must be all-gathered before the contiguous window + can be sliced out (``data_is_cp_sharded=True``). + * Megatron worker: only the model input and the returned logits are + CP-sharded; the microbatch dict the loss sees still holds the FULL + ``[B, S]`` tensors, so this rank's window is a plain ``narrow`` + (``data_is_cp_sharded=False``). + + No-op without a CP group (or at CP world size 1), so both single-GPU paths + are byte-identical to before. + """ + if cp_group is None or torch.distributed.get_world_size(cp_group) <= 1: + return x + if data_is_cp_sharded: + return cp_load_balanced_to_contiguous(x, cp_group=cp_group, seq_dim=seq_dim) + local = to_local_if_dtensor(x) + cp_size = torch.distributed.get_world_size(cp_group) + cp_rank = torch.distributed.get_rank(cp_group) + local_len = local.shape[seq_dim] // cp_size + return local.narrow(seq_dim, cp_rank * local_len, local_len).contiguous() + + +def prepare_csd_cross_tokenizer_loss_input( + logits: torch.Tensor, + data: Mapping[str, Any], + *, + projection_matrix_paths: list[Optional[str]], + cross_tokenizer_teacher_mask: Optional[list[bool]] = None, + vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, + context_parallel_group: Optional[torch.distributed.ProcessGroup] = None, + cp_sharder: Optional["ContextParallelSharder"] = None, +) -> tuple[ + torch.Tensor, + Dict[int, torch.Tensor], + Dict[int, SparseTeacherLogits], + Dict[int, LocalizedAlignment], + Dict[int, int], + Optional[torch.distributed.ProcessGroup], + Optional[torch.distributed.ProcessGroup], + Optional[torch.distributed.ProcessGroup], +]: + """Build the per-teacher cross-tokenizer distillation loss pieces from student logits + IPC teacher data. + + Rebuilds each teacher's dense full-vocab or sparse top-k + logZ logits from + its per-rank CUDA IPC handles and does the shared CP-resolution the loss + needs. The contiguous student logits / input_ids / token_mask are relaid + once and shared across teachers. Per teacher, a localized alignment is built: a + cross-tokenizer teacher (``projection_matrix_paths[i]`` set) gets the + localized, next-token-shifted chunk alignment from its ``alignment_{i}_*`` + keys; a same-tokenizer teacher (``None`` path) gets a thin alignment carrying + only the shared student fields (identity 1:1 token alignment, no chunks). + The TP group comes from the student ``logits``' device mesh, falling back to + ``vocab_parallel_group`` for non-DTensor logits; the CP group is always the + passed ``context_parallel_group``. + + Args: + projection_matrix_paths: Per-teacher projection paths. Its length is the + teacher count and drives the ``teacher_{i}_*`` / ``alignment_{i}_*`` + keys read here; a ``None`` entry marks a same-tokenizer teacher. + cp_sharder: Automodel's model-owned sequence layout, set only when + context_parallel_size > 1. It relays the student tensors to this + rank's contiguous window. + + Returns: + ``(student_logits_contig, teacher_full_logits_by_idx, + teacher_sparse_logits_by_idx, aligns_by_idx, + dense_reconstruction_fallbacks_by_idx, tp_group, cp_group, + dp_cp_group)``. The reconstruction counts come from the actual dense + rebuild branch. ``dp_cp_group`` is the group the loss reduces its global + normalizers over; see :func:`loss_replica_group` for why it must not be + ``WORLD`` under pipeline parallelism. + """ + if isinstance(logits, DTensor): + mesh = logits.device_mesh + mesh_names = mesh.mesh_dim_names or () + tp_group = mesh.get_group("tp") if "tp" in mesh_names else None + else: + tp_group = vocab_parallel_group + cp_group = context_parallel_group + data_is_cp_sharded = False + + device = torch.cuda.current_device() + + # Student CP relay is computed once and shared by every teacher's KD term + # and the next-token-accuracy metric. + # Automodel restores its own model layout before NeMo RL selects the + # contiguous IPC-consumer window. + if cp_sharder is not None: + # Keep the student-logit autograd graph intact: to_local_if_dtensor() + # runs DTensor.to_local() under torch.no_grad(), which would detach the + # CSD KD loss from the student model for TP > 1. + local_logits = logits.to_local() if isinstance(logits, DTensor) else logits + full_student_logits = cp_sharder.gather_token_tensor( + local_logits, seq_dim=1, trim=True + ) + cp_size = ( + torch.distributed.get_world_size(cp_group) if cp_group is not None else 1 + ) + full_student_seq_len = full_student_logits.shape[1] + if full_student_seq_len % cp_size != 0: + raise ValueError( + "CSD student sequence length must be divisible by the student " + "context parallel size, but got " + f"sequence_length={full_student_seq_len}, cp_size={cp_size}. " + "Set policy.make_sequence_length_divisible_by to a multiple of " + "policy.automodel_cfg.context_parallel_size." + ) + cp_rank = torch.distributed.get_rank(cp_group) if cp_group is not None else 0 + student_seq_len = full_student_seq_len // cp_size + student_seq_start = cp_rank * student_seq_len + student_logits_contig = full_student_logits.narrow( + 1, student_seq_start, student_seq_len + ).contiguous() + student_input_ids = to_local_if_dtensor(data["input_ids"])[ + :, student_seq_start : student_seq_start + student_seq_len + ].contiguous() + student_token_mask = to_local_if_dtensor(data["token_mask"])[ + :, student_seq_start : student_seq_start + student_seq_len + ].contiguous() + student_kd_token_mask = to_local_if_dtensor( + data.get("kd_token_mask", data["token_mask"]) + )[:, student_seq_start : student_seq_start + student_seq_len].contiguous() + else: + # Logits are load-balanced on both backends. Student-sequence data is + # load-balanced only for the legacy DTensor path; Megatron leaves it + # full, so select its contiguous window directly. + student_logits_contig = cp_load_balanced_to_contiguous( + logits, cp_group=cp_group + ) + student_input_ids = _student_seq_to_contiguous_window( + data["input_ids"], cp_group, data_is_cp_sharded=data_is_cp_sharded + ) + student_token_mask = _student_seq_to_contiguous_window( + data["token_mask"], cp_group, data_is_cp_sharded=data_is_cp_sharded + ) + student_kd_token_mask = _student_seq_to_contiguous_window( + data.get("kd_token_mask", data["token_mask"]), + cp_group, + data_is_cp_sharded=data_is_cp_sharded, + ) + sample_mask = to_local_if_dtensor(data["sample_mask"]) + + teacher_full_logits_by_idx: Dict[int, torch.Tensor] = {} + teacher_sparse_logits_by_idx: Dict[int, SparseTeacherLogits] = {} + aligns_by_idx: Dict[int, LocalizedAlignment] = {} + dense_reconstruction_fallbacks_by_idx: Dict[int, int] = {} + student_cp_size = ( + torch.distributed.get_world_size(cp_group) if cp_group is not None else 1 + ) + if cross_tokenizer_teacher_mask is None: + cross_tokenizer_teacher_mask = [p is not None for p in projection_matrix_paths] + if len(cross_tokenizer_teacher_mask) != len(projection_matrix_paths): + raise ValueError("cross_tokenizer_teacher_mask must have one entry per teacher") + for i, is_cross_tokenizer in enumerate(cross_tokenizer_teacher_mask): + sparse_key = f"teacher_{i}_sparse_logits_ipc" + full_key = f"teacher_{i}_full_logits_ipc" + has_sparse_logits = sparse_key in data + has_full_logits = full_key in data + if has_sparse_logits == has_full_logits: + raise ValueError( + f"Teacher {i} must provide exactly one dense or sparse logits " + f"IPC payload; dense={has_full_logits}, sparse={has_sparse_logits}." + ) + if has_sparse_logits: + if not is_cross_tokenizer: + raise ValueError(f"Same-vocab teacher {i} cannot use sparse CSD IPC.") + teacher_sparse_logits = rebuild_teacher_sparse_logits_from_ipc( + data[sparse_key], + device=device, + ) + teacher_sparse_logits_by_idx[i] = teacher_sparse_logits + full_teacher_seq_len = int(teacher_sparse_logits[0].shape[1]) + if full_teacher_seq_len % student_cp_size != 0: + raise ValueError( + "Sparse teacher sequence length must be divisible by the " + f"student CP size; teacher={i}, seq={full_teacher_seq_len}, " + f"student_cp={student_cp_size}." + ) + teacher_seq_len = full_teacher_seq_len // student_cp_size + else: + ( + teacher_full_logits, + dense_reconstruction_fallbacks, + ) = rebuild_teacher_full_logits_from_ipc( + data[full_key], + cp_group=cp_group, + device=device, + ) + teacher_full_logits_by_idx[i] = teacher_full_logits + dense_reconstruction_fallbacks_by_idx[i] = dense_reconstruction_fallbacks + teacher_seq_len = int(teacher_full_logits.shape[1]) + if not is_cross_tokenizer: + # Same-tokenizer teacher: identity token alignment, no chunk + # localization. Carry only the shared student fields. + aligns_by_idx[i] = LocalizedAlignment( + sample_mask=sample_mask, + student_input_ids=student_input_ids, + student_token_mask=student_token_mask, + student_kd_token_mask=student_kd_token_mask, + ) + continue + alignment_prefix = f"alignment_{i}_" + if cp_sharder is not None: + student_chunk_id_source_full = to_local_if_dtensor( + data[f"{alignment_prefix}student_chunk_id"] + ) + student_chunk_id_contig = student_chunk_id_source_full[ + :, student_seq_start : student_seq_start + student_seq_len + ].contiguous() + + teacher_chunk_id_source_full = to_local_if_dtensor( + data[f"{alignment_prefix}teacher_chunk_id"] + ) + teacher_seq_start = cp_rank * teacher_seq_len + teacher_chunk_id_contig = teacher_chunk_id_source_full[ + :, teacher_seq_start : teacher_seq_start + teacher_seq_len + ].contiguous() + align = LocalizedAlignment( + sample_mask=sample_mask, + pair_valid=to_local_if_dtensor(data[f"{alignment_prefix}pair_valid"]), + pair_is_correct=to_local_if_dtensor( + data[f"{alignment_prefix}pair_is_correct"] + ), + ) + else: + align = localize_alignment( + data, + teacher_seq_len=teacher_seq_len, + alignment_prefix=alignment_prefix, + ) + student_chunk_id_contig = _student_seq_to_contiguous_window( + align.student_chunk_id, + cp_group, + data_is_cp_sharded=data_is_cp_sharded, + ) + teacher_cp_rank = ( + torch.distributed.get_rank(cp_group) if cp_group is not None else 0 + ) + teacher_chunk_id_contig = to_local_if_dtensor( + data[f"{alignment_prefix}teacher_chunk_id"] + )[ + :, + teacher_cp_rank * teacher_seq_len : (teacher_cp_rank + 1) + * teacher_seq_len, + ] + cp_rank = ( + torch.distributed.get_rank(cp_group) + if cp_group is not None + and torch.distributed.get_world_size(cp_group) > 1 + else 0 + ) + teacher_seq_start = cp_rank * teacher_seq_len + + # Contiguous, UNSHIFTED chunk ids -> per-chunk position spans for the CSD + # path. CSD reads spans and applies its own per-chunk kl_chunk_shift, so + # it must NOT see the P-KL/gold next-token-shifted chunk ids. + max_pairs = align.pair_valid.shape[1] + # GLOBAL spans: the CSD KD term gathers the student/teacher logits to the + # full sequence, so its spans must index global positions. Gather the + # contiguous-window chunk ids to the full sequence before deriving spans + # (no-op at CP=1). + if cp_sharder is not None: + # Automodel has already restored these data tensors to the full, + # canonical sequence before the local contiguous window is sliced. + student_chunk_id_global = student_chunk_id_source_full + teacher_chunk_id_global = teacher_chunk_id_source_full + else: + student_chunk_id_global = allgather_cp_contiguous_tensor( + student_chunk_id_contig, cp_group + ) + teacher_chunk_id_global = allgather_cp_contiguous_tensor( + teacher_chunk_id_contig, cp_group + ) + align.student_spans = _chunk_ids_to_spans(student_chunk_id_global, max_pairs) + align.teacher_spans = _chunk_ids_to_spans(teacher_chunk_id_global, max_pairs) + # Preserve the aligner's real per-sample chunk count rather than + # iterating every padded pair slot; pair_valid remains the final gate. + align.num_chunks = to_local_if_dtensor(align.pair_valid).shape[1] * torch.ones( + align.pair_valid.shape[0], dtype=torch.long, device=align.pair_valid.device + ) + # Teacher input ids for this CP rank's contiguous teacher window (matches + # the teacher-logit / teacher_chunk_id slice). + teacher_ids_full = to_local_if_dtensor(data[f"teacher_{i}_input_ids"]) + align.teacher_input_ids = teacher_ids_full[ + :, teacher_seq_start : teacher_seq_start + teacher_seq_len + ] + # Preserve the established next-token-shifted localized fields for + # metrics and compatibility while CSD consumes the unshifted spans. + if cp_sharder is not None: + student_chunk_id_shifted = student_chunk_id_source_full.roll( + shifts=-1, dims=1 + ) + student_chunk_id_shifted[:, -1] = -1 + align.student_chunk_id = student_chunk_id_shifted[ + :, student_seq_start : student_seq_start + student_seq_len + ].contiguous() + teacher_chunk_id_shifted = teacher_chunk_id_source_full.roll( + shifts=-1, dims=1 + ) + teacher_chunk_id_shifted[:, -1] = -1 + align.teacher_chunk_id = teacher_chunk_id_shifted[ + :, teacher_seq_start : teacher_seq_start + teacher_seq_len + ].contiguous() + else: + align.student_chunk_id = cp_shift_next( + student_chunk_id_contig, cp_group, fill=-1 + ) + align.teacher_chunk_id = cp_shift_next( + teacher_chunk_id_contig, cp_group, fill=-1 + ) + align.student_input_ids = student_input_ids + align.student_token_mask = student_token_mask + align.student_kd_token_mask = student_kd_token_mask + aligns_by_idx[i] = align + return ( + student_logits_contig, + teacher_full_logits_by_idx, + teacher_sparse_logits_by_idx, + aligns_by_idx, + dense_reconstruction_fallbacks_by_idx, + tp_group, + cp_group, + loss_replica_group(cp_group), + ) + + +def prepare_csd_owner_local_loss_input( + logits: torch.Tensor, + data: Mapping[str, Any], + *, + cross_tokenizer_teacher_mask: list[bool], + vocab_parallel_group: Optional[torch.distributed.ProcessGroup], + context_parallel_group: Optional[torch.distributed.ProcessGroup], +) -> tuple[ + Optional[torch.Tensor], + Dict[int, torch.Tensor], + Dict[int, TeacherRowPayload], + Dict[int, LocalizedAlignment], + Optional[torch.distributed.ProcessGroup], + Optional[torch.distributed.ProcessGroup], + Optional[torch.distributed.ProcessGroup], +]: + """Build CSD loss pieces that keep the Megatron student logits in their CP shard. + + Cross-tokenizer teachers are not reassembled: each gets a + :class:`TeacherRowPayload` holding its IPC records, so the loss reads only + the teacher rows of the chunks this CP rank owns. Their alignments carry + global spans, the full student/teacher ids and the global position of every + local logit row. The student logits are never gathered across CP, except + for a same-tokenizer teacher, which keeps today's contiguous-window + reconstruction (the window is built only when such a teacher exists). + + Args: + logits: Megatron student logits ``[B, S_pad/CP, V/TP]`` in the + load-balanced CP layout. + data: Microbatch dict holding full ``[B, S]`` student tensors. + cross_tokenizer_teacher_mask: Per-teacher cross-tokenizer flags. + vocab_parallel_group: Student tensor-parallel group. + context_parallel_group: Student context-parallel group. + + Returns: + ``(student_logits_contig, teacher_full_logits_by_idx, + teacher_row_payloads_by_idx, aligns_by_idx, tp_group, cp_group, + dp_cp_group)``. ``student_logits_contig`` and + ``teacher_full_logits_by_idx`` are set only for same-tokenizer teachers. + """ + if isinstance(logits, DTensor): + raise NotImplementedError( + "Owner-local CSD evaluation supports Megatron students only." + ) + tp_group = vocab_parallel_group + cp_group = context_parallel_group + device = logits.device + student_positions = cp_load_balanced_global_positions( + int(logits.shape[1]), cp_group, device=device + ) + full_student_ids = to_local_if_dtensor(data["input_ids"]) + cp_size = ( + torch.distributed.get_world_size(cp_group) + if cp_group is not None and torch.distributed.is_initialized() + else 1 + ) + padded_seq_len = int(logits.shape[1]) * cp_size + if full_student_ids.shape[1] > padded_seq_len: + raise ValueError( + "Owner-local CSD student logits cover fewer positions than the " + f"microbatch: {padded_seq_len} logit positions for " + f"{full_student_ids.shape[1]} input ids." + ) + full_token_mask = to_local_if_dtensor(data["token_mask"]) + full_kd_token_mask = to_local_if_dtensor( + data.get("kd_token_mask", data["token_mask"]) + ) + sample_mask = to_local_if_dtensor(data["sample_mask"]) + + student_logits_contig: Optional[torch.Tensor] = None + contiguous_fields: Optional[tuple[torch.Tensor, ...]] = None + teacher_full_logits_by_idx: Dict[int, torch.Tensor] = {} + teacher_row_payloads_by_idx: Dict[int, TeacherRowPayload] = {} + aligns_by_idx: Dict[int, LocalizedAlignment] = {} + for i, is_cross_tokenizer in enumerate(cross_tokenizer_teacher_mask): + sparse_key = f"teacher_{i}_sparse_logits_ipc" + full_key = f"teacher_{i}_full_logits_ipc" + has_sparse_logits = sparse_key in data + if has_sparse_logits == (full_key in data): + raise ValueError( + f"Teacher {i} must provide exactly one dense or sparse logits " + f"IPC payload; dense={full_key in data}, sparse={has_sparse_logits}." + ) + if not is_cross_tokenizer: + if has_sparse_logits: + raise ValueError(f"Same-vocab teacher {i} cannot use sparse CSD IPC.") + if contiguous_fields is None: + student_logits_contig = cp_load_balanced_to_contiguous( + logits, cp_group=cp_group + ) + contiguous_fields = tuple( + _student_seq_to_contiguous_window( + value, cp_group, data_is_cp_sharded=False + ) + for value in (full_student_ids, full_token_mask, full_kd_token_mask) + ) + teacher_full_logits_by_idx[i], _ = rebuild_teacher_full_logits_from_ipc( + data[full_key], cp_group=cp_group, device=logits.get_device() + ) + aligns_by_idx[i] = LocalizedAlignment( + sample_mask=sample_mask, + student_input_ids=contiguous_fields[0], + student_token_mask=contiguous_fields[1], + student_kd_token_mask=contiguous_fields[2], + ) + continue + + teacher_row_payloads_by_idx[i] = TeacherRowPayload( + transport="sparse" if has_sparse_logits else "dense", + entries=list(data[sparse_key if has_sparse_logits else full_key]), + ) + pair_valid = to_local_if_dtensor(data[f"alignment_{i}_pair_valid"]) + max_pairs = int(pair_valid.shape[1]) + aligns_by_idx[i] = LocalizedAlignment( + sample_mask=sample_mask, + pair_valid=pair_valid, + pair_is_correct=to_local_if_dtensor(data[f"alignment_{i}_pair_is_correct"]), + student_input_ids=full_student_ids, + student_token_mask=full_token_mask, + student_kd_token_mask=full_kd_token_mask, + teacher_input_ids=to_local_if_dtensor(data[f"teacher_{i}_input_ids"]), + # The Megatron microbatch holds full alignment tensors, so global + # spans need no sequence collective. + student_spans=_chunk_ids_to_spans( + to_local_if_dtensor(data[f"alignment_{i}_student_chunk_id"]), + max_pairs, + ), + teacher_spans=_chunk_ids_to_spans( + to_local_if_dtensor(data[f"alignment_{i}_teacher_chunk_id"]), + max_pairs, + ), + num_chunks=torch.full( + (pair_valid.shape[0],), + max_pairs, + dtype=torch.long, + device=pair_valid.device, + ), + student_global_positions=student_positions, + ) + return ( + student_logits_contig, + teacher_full_logits_by_idx, + teacher_row_payloads_by_idx, + aligns_by_idx, + tp_group, + cp_group, + loss_replica_group(cp_group), + ) + + def prepare_xtoken_cross_tokenizer_loss_input( logits: torch.Tensor, data: Mapping[str, Any], @@ -1266,7 +3081,7 @@ def prepare_xtoken_cross_tokenizer_loss_input( teacher_full_logits_by_idx: Dict[int, torch.Tensor] = {} aligns_by_idx: Dict[int, LocalizedAlignment] = {} for i, proj_path in enumerate(projection_matrix_paths): - teacher_full_logits = rebuild_teacher_full_logits_from_ipc( + teacher_full_logits, _ = rebuild_teacher_full_logits_from_ipc( data[f"teacher_{i}_full_logits_ipc"], cp_group=cp_group, device=device, diff --git a/nemo_rl/algorithms/x_token/token_aligner.py b/nemo_rl/algorithms/x_token/token_aligner.py index 13b35e54755..de6130e13d3 100644 --- a/nemo_rl/algorithms/x_token/token_aligner.py +++ b/nemo_rl/algorithms/x_token/token_aligner.py @@ -252,7 +252,8 @@ class AlignmentBatch: V2 can shard on dim 0 without knowing about cross-tokenizer specifics. Attributes: - pair_valid: ``[B, max_pairs]`` bool. False on padding entries. + pair_valid: ``[B, max_pairs]`` bool. False on padding entries and on + pairs whose student or teacher tokens are all tokenizer padding. pair_is_correct: ``[B, max_pairs]`` bool. True when decoded student and teacher span texts are equal, when an NFC-validated native repair matches them, or when both spans are their tokenizer's @@ -294,7 +295,7 @@ def __init__( self, student_tokenizer, teacher_tokenizer, - projection_matrix_path: str, + projection_matrix_path: str | None, ): self.student_tokenizer = student_tokenizer self.teacher_tokenizer = teacher_tokenizer @@ -533,22 +534,41 @@ def _drop_padding( student_attention_mask: torch.Tensor | None, teacher_attention_mask: torch.Tensor | None, ) -> None: - """Strip tokenizer padding out of the chunk-id tensors. + """Strip tokenizer padding out of the chunk ids and invalidate padding pairs. Mutates ``batch`` in place. For every position the attention mask marks as padding, reset ``*_chunk_id`` to ``-1``. Gating per position (rather than trimming a contiguous span) keeps this correct under either - left- or right-padding. A pair whose tokens are entirely padding on - one side then has size 0 there and is dropped by - :func:`nemo_rl.algorithms.x_token.loss_utils.valid_chunk_mask`; a pair - straddling the real/pad boundary shrinks to its real tokens. + left- or right-padding; a pair straddling the real/pad boundary shrinks + to its real tokens. When both masks are given, a pair with tokens on + both sides whose student or teacher tokens are all padding becomes + invalid in ``pair_valid``. A pair with no tokens on one side (one-sided + by alignment) keeps its flag. """ + max_pairs = batch.pair_valid.shape[1] + two_sided = TokenAligner._pairs_with_tokens( + batch.student_chunk_id, max_pairs + ) & TokenAligner._pairs_with_tokens(batch.teacher_chunk_id, max_pairs) if student_attention_mask is not None: s_pad = student_attention_mask == 0 batch.student_chunk_id[s_pad] = -1 if teacher_attention_mask is not None: t_pad = teacher_attention_mask == 0 batch.teacher_chunk_id[t_pad] = -1 + if student_attention_mask is None or teacher_attention_mask is None: + return + two_sided_real = TokenAligner._pairs_with_tokens( + batch.student_chunk_id, max_pairs + ) & TokenAligner._pairs_with_tokens(batch.teacher_chunk_id, max_pairs) + batch.pair_valid &= two_sided_real | ~two_sided + + @staticmethod + def _pairs_with_tokens(chunk_id: torch.Tensor, max_pairs: int) -> torch.Tensor: + """``[B, max_pairs]`` bool: True where the pair owns a position of ``chunk_id``.""" + # Column 0 counts the -1 sentinel; column p + 1 counts pair p. + counts = torch.zeros((chunk_id.shape[0], max_pairs + 1), dtype=torch.long) + counts.scatter_add_(1, chunk_id + 1, torch.ones_like(chunk_id)) + return counts[:, 1:] > 0 @staticmethod def _pairs_to_batch( diff --git a/nemo_rl/algorithms/x_token/utils.py b/nemo_rl/algorithms/x_token/utils.py index dcb2acfb3e5..e99e7542fd0 100644 --- a/nemo_rl/algorithms/x_token/utils.py +++ b/nemo_rl/algorithms/x_token/utils.py @@ -84,6 +84,8 @@ def assert_xtoken_ipc_node_local( teacher_cp: int, student_dp: int, teacher_dp: int, + student_pp: int, + teacher_pp: int, ) -> None: """Fail fast if the teacher->student logit transport would cross a node. @@ -92,12 +94,16 @@ def assert_xtoken_ipc_node_local( A single-node job is therefore always safe; a multi-node job is only safe when every student rank's required teacher shards live on its own node. + The model-parallel group is ``tp * cp * pp``. Pipeline parallelism counts + even though only the last stage produces logits: the group must stay inside + one node for that stage to be co-located with the student ranks that import + its buffers. """ if num_nodes <= 1: return - student_group = student_tp * student_cp - teacher_group = teacher_tp * teacher_cp + student_group = student_tp * student_cp * student_pp + teacher_group = teacher_tp * teacher_cp * teacher_pp assert teacher_dp == student_dp, ( "Multi-node xtoken distillation uses node-local CUDA IPC for the teacher " "logits, which requires teacher and student to share a data_parallel " @@ -107,18 +113,18 @@ def assert_xtoken_ipc_node_local( ) assert teacher_group == student_group, ( "Multi-node xtoken distillation requires teacher and student to share " - f"the same model-parallel group size tp*cp (got teacher={teacher_group}, " + f"the same model-parallel group size tp*cp*pp (got teacher={teacher_group}, " f"student={student_group}); differing sizes imply a non-colocated grid " "whose teacher shards are not node-local." ) assert student_group <= gpus_per_node, ( - f"Multi-node xtoken distillation needs the model-parallel group tp*cp " + f"Multi-node xtoken distillation needs the model-parallel group tp*cp*pp " f"({student_group}) to fit within one node ({gpus_per_node} GPUs); a " "group spanning nodes makes the teacher-logit IPC cross-node." ) assert gpus_per_node % student_group == 0, ( f"Multi-node xtoken distillation needs gpus_per_node ({gpus_per_node}) " - f"to be a multiple of the model-parallel group tp*cp ({student_group}) " + f"to be a multiple of the model-parallel group tp*cp*pp ({student_group}) " "so DP groups are node-aligned and never straddle a node boundary." ) diff --git a/nemo_rl/algorithms/xtoken_off_policy_distillation.py b/nemo_rl/algorithms/xtoken_off_policy_distillation.py index 6856108e997..ea96b668985 100644 --- a/nemo_rl/algorithms/xtoken_off_policy_distillation.py +++ b/nemo_rl/algorithms/xtoken_off_policy_distillation.py @@ -18,15 +18,14 @@ rollout, no generation). Per step: 1. Pull a collated batch (student & teacher token ids + alignment). - 2. Run each teacher forward via ``Policy.get_full_logits_ipc`` on TEACHER - token ids — every teacher ships full-vocab logits over CUDA IPC (no - driver round-trip), reassembled across the teacher's TP/CP shards on - the student side. + 2. Run each teacher forward on TEACHER token ids and ship either dense + full-vocabulary or sparse top-k logits over CUDA IPC (no driver + round-trip), reassembled across the teacher's TP/CP shards on the + student side. 3. Pack alignment payload + teacher IPC handles into a student-side ``train_data`` dict. 4. ``student_policy.train(train_data, loss_fn)`` — student forward + - loss + backward + optimizer step happens inside the dtensor v2 - worker. + loss + backward + optimizer step happens inside the policy worker. The collator and aligner do all the CPU-side cross-tokenizer work; the loss function does only loss math; this module is just plumbing. @@ -34,22 +33,34 @@ from __future__ import annotations +import hashlib +import json import math import os -from typing import Any, NotRequired, Optional, TypedDict, cast +import re +from collections.abc import Iterator, Mapping, Sequence +from functools import partial +from itertools import count +from typing import Any, Literal, NotRequired, Optional, TypedDict, cast import numpy as np import torch -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator from torchdata.stateful_dataloader import StatefulDataLoader from transformers.tokenization_utils_base import PreTrainedTokenizerBase +from nemo_rl.algorithms.loss.csd import ( + CSDLossConfig, + CSDLossFn, + owner_local_config_violations, +) from nemo_rl.algorithms.loss.loss_functions import ( CrossTokenizerDistillationLossConfig, CrossTokenizerDistillationLossFn, ) from nemo_rl.algorithms.utils import set_seed from nemo_rl.algorithms.x_token import TokenAligner +from nemo_rl.algorithms.x_token.loss_utils import _chunk_ids_to_spans from nemo_rl.algorithms.x_token.utils import ( assert_teacher_student_batch_grid, assert_xtoken_ipc_node_local, @@ -60,10 +71,18 @@ CrossTokenizerCollator, CrossTokenizerCollatorConfig, ) +from nemo_rl.data.cross_tokenizer.csd_collate import CSDCrossTokenizerCollator from nemo_rl.data.datasets import AllTaskProcessedDataset +from nemo_rl.data.packing import ( + LockstepPackingItem, + LockstepPackingPlan, + SidePackingSpec, + build_lockstep_packing_plan, +) from nemo_rl.data.utils import load_dataloader_state from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster +from nemo_rl.models.megatron.router_replay import router_replay_enabled from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.utils.checkpoint import ( @@ -81,7 +100,7 @@ # pre-flight (which assumes [B, student_seq, ...] for every 2+D tensor) must # skip them. Sources: # - teacher_full_logits_ipc: list[B] of CUDA IPC handle dicts produced by -# FullLogitsPostProcessor in automodel_policy_worker.get_full_logits_ipc. +# FullLogitsPostProcessor in the Automodel/Megatron worker's get_full_logits_ipc. # Not a tensor at all — list of dicts — but listed here so the worker's # dict-level dim check skips it. # - teacher_input_ids/teacher_token_mask + alignment_*: produced by @@ -97,9 +116,11 @@ def xtoken_non_student_seq_keys( per microbatch, but the worker's ``check_sequence_dim`` pre-flight (which assumes ``[B, student_seq, ...]`` for every 2+D tensor) must skip them. The set is teacher-dependent, so it is built from the loss fn's per-teacher - metadata rather than a static constant. Every teacher ships full logits, so - each contributes ``teacher_{i}_full_logits_ipc`` (a list of CUDA IPC handle - dicts, not a tensor); a cross-tokenizer teacher additionally rides its + metadata rather than a static constant. Every teacher contributes exactly + one of ``teacher_{i}_full_logits_ipc`` or + ``teacher_{i}_sparse_logits_ipc`` (lists of CUDA IPC handle dicts, not + tensors); both possible keys are skipped so the transport may be selected + at runtime. A cross-tokenizer teacher additionally rides its teacher-seq tokenization (``teacher_{i}_input_ids`` / ``teacher_{i}_token_mask``, ``[B, T_t]``) and its teacher-seq / max_pairs ``alignment_{i}_*`` keys. The ``alignment_{i}_student_chunk_id`` (``[B, T_s]``) follows the student-seq @@ -108,7 +129,9 @@ def xtoken_non_student_seq_keys( keys: set[str] = set() for i in range(loss_fn.num_teachers): keys.add(f"teacher_{i}_full_logits_ipc") - if loss_fn.projection_matrix_paths[i] is not None: # cross-tokenizer + if isinstance(loss_fn, CSDLossFn): + keys.add(f"teacher_{i}_sparse_logits_ipc") + if _loss_teacher_is_cross_tokenizer(loss_fn, i): # cross-tokenizer keys.add(f"teacher_{i}_input_ids") keys.add(f"teacher_{i}_token_mask") keys.add(f"alignment_{i}_pair_valid") @@ -117,6 +140,306 @@ def xtoken_non_student_seq_keys( return frozenset(keys) +# Per-microbatch metrics that are *intensive* — rates, ratios, per-chunk +# averages, and global-batch constants replicated onto every microbatch. Their +# per-step value is the mean over the flat per-(microbatch, DP rank) list; +# everything else is an *extensive* per-microbatch contribution (loss shares +# normalized by a global denominator, counts) whose per-step value is the sum. +MEAN_REDUCED_MB_METRICS = frozenset( + { + "lr", + "wd", + "global_valid_seqs", + "global_valid_toks", + "accuracy", + "proj_accuracy", + "kl_loss_scale", + # Gold-loss path: per-mb local chunk average. + "kl_common", + "l1_uncommon", + # CSD diagnostics. Each microbatch + # reports an average over *its own* chunks, so summing them scales + # with the microbatch count instead of staying per-chunk (which is + # how top1_acc_per_chunk reached an impossible 110.60). + "kl_common_per_chunk", + "kl_partition_last_per_chunk", + "top1_acc_per_chunk", + } +) + +# Per-teacher metrics are suffixed ``_t{i}`` by the loss fn's aggregation. +_TEACHER_METRIC_SUFFIX = re.compile(r"_t\d+$") + +_CSD_LOCKSTEP_PACKING_ALGORITHM = "lockstep_first_fit_decreasing" +_UNSAFE_PACKED_MODEL_MARKERS = ( + "mamba", + "nemotron-h", + "nemotron_h", + "nemotronh", + "nano", +) + +_CSD_LOGICAL_BATCH_DIGEST_SCHEMA_VERSION = 1 +_CSD_LOGICAL_BATCH_BASE_FIELDS = ( + "input_ids", + "input_lengths", + "token_mask", + "kd_token_mask", + "sample_mask", + "student_semantic_regions", +) +_CSD_LOGICAL_BATCH_TEACHER_SUFFIXES = ( + "_input_ids", + "_input_lengths", + "_token_mask", + "_semantic_regions", +) + + +def reduce_mb_metric(key: str, values: Any) -> float: + """Reduce one metric's per-(microbatch, DP rank) list to its per-step scalar.""" + if _TEACHER_METRIC_SUFFIX.sub("", key) in MEAN_REDUCED_MB_METRICS: + return float(np.mean(values)) + return float(np.sum(values)) + + +def _log_dense_ipc_reconstruction_telemetry( + all_mb_metrics: Mapping[str, Any], + *, + packing_plan: LockstepPackingPlan, + num_teachers: int, +) -> None: + """Validate and log consumer-measured packed dense-IPC fallbacks.""" + fallback_metric_keys = ( + "ipc_reconstruction_fallbacks", + *(f"ipc_reconstruction_fallbacks_t{i}" for i in range(num_teachers)), + ) + missing_fallback_metrics = [ + key for key in fallback_metric_keys if key not in all_mb_metrics + ] + if missing_fallback_metrics: + raise RuntimeError( + "Packed dense IPC reconstruction telemetry is missing consumer-side " + f"metrics {missing_fallback_metrics}." + ) + + reduced_fallbacks = { + key: reduce_mb_metric(key, all_mb_metrics[key]) for key in fallback_metric_keys + } + non_integral_fallbacks = { + key: value + for key, value in reduced_fallbacks.items() + if not math.isfinite(value) + or value < 0 + or not math.isclose(value, round(value), abs_tol=1.0e-6) + } + if non_integral_fallbacks: + raise RuntimeError( + "Packed dense IPC reconstruction telemetry must contain finite, " + "non-negative integer counts after reduction; got " + f"{non_integral_fallbacks}." + ) + fallback_count = round(reduced_fallbacks["ipc_reconstruction_fallbacks"]) + per_teacher_fallback_counts = tuple( + round(reduced_fallbacks[f"ipc_reconstruction_fallbacks_t{i}"]) + for i in range(num_teachers) + ) + if fallback_count != sum(per_teacher_fallback_counts): + raise RuntimeError( + "Packed dense IPC reconstruction telemetry is internally " + f"inconsistent: total={fallback_count}, " + f"per_teacher={per_teacher_fallback_counts}." + ) + teacher_counts = " ".join( + f"teacher_{i}={count}" for i, count in enumerate(per_teacher_fallback_counts) + ) + print( + "CSD_IPC_RECONSTRUCTION " + f"batch_uid={packing_plan.batch_uid} " + f"reconstruction_fallbacks={fallback_count} " + f"{teacher_counts}", + flush=True, + ) + + +def _canonical_logical_digest_json_value(value: Any, *, field: str) -> Any: + """Normalize one non-tensor logical field for deterministic JSON hashing.""" + if isinstance(value, np.generic): + return _canonical_logical_digest_json_value(value.item(), field=field) + if value is None or isinstance(value, (bool, int, float, str)): + if isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"CSD logical digest field {field!r} contains {value!r}.") + return value + if isinstance(value, Mapping): + normalized: dict[str, Any] = {} + non_string_keys = [key for key in value if not isinstance(key, str)] + if non_string_keys: + raise TypeError( + f"CSD logical digest field {field!r} has non-string " + f"mapping keys {non_string_keys!r}." + ) + for key in sorted(value): + normalized[key] = _canonical_logical_digest_json_value( + value[key], field=f"{field}.{key}" + ) + return normalized + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return [ + _canonical_logical_digest_json_value(item, field=f"{field}[{index}]") + for index, item in enumerate(value) + ] + raise TypeError( + f"CSD logical digest field {field!r} has unsupported type " + f"{type(value).__name__}." + ) + + +def _logical_digest_json_scalar_count(value: Any) -> int: + """Count canonical JSON scalar leaves for a ragged logical field.""" + if isinstance(value, Mapping): + return sum(_logical_digest_json_scalar_count(item) for item in value.values()) + if isinstance(value, list): + return sum(_logical_digest_json_scalar_count(item) for item in value) + return 1 + + +def _logical_digest_field_record(value: Any, *, field: str) -> dict[str, Any]: + """Return shape, dtype, count, and a framed SHA-256 for one field.""" + if torch.is_tensor(value): + tensor = value.detach().cpu().contiguous() + metadata = { + "shape": list(tensor.shape), + "dtype": str(tensor.dtype), + "count": int(tensor.numel()), + } + header = json.dumps( + metadata, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + raw = tensor.view(torch.uint8).numpy().tobytes(order="C") + digest = hashlib.sha256(header + b"\n" + raw).hexdigest() + return {**metadata, "sha256": digest} + + normalized = _canonical_logical_digest_json_value(value, field=field) + shape = [len(normalized)] if isinstance(normalized, list) else [] + metadata = { + "shape": shape, + "dtype": "canonical-json", + "count": _logical_digest_json_scalar_count(normalized), + } + payload = json.dumps( + {"metadata": metadata, "value": normalized}, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return {**metadata, "sha256": hashlib.sha256(payload).hexdigest()} + + +def build_csd_logical_batch_digest_record( + batch: BatchedDataDict[Any], *, batch_uid: int +) -> dict[str, Any]: + """Build the arm-independent logical batch record logged before packing. + + The record intentionally contains hashes and structural metadata only. The + frozen asset keeps the ordered sample IDs beside this record, while runtime + logs can prove exact equality without exposing dataset text or identifiers. + """ + if batch_uid < 0 or batch_uid >= 2**31: + raise ValueError( + f"batch_uid must fit the non-negative high 31 bits of int64, got {batch_uid}." + ) + sample_ids_value = batch.get("sample_id") + if torch.is_tensor(sample_ids_value): + sample_ids: Any = sample_ids_value.detach().cpu().tolist() + elif isinstance(sample_ids_value, Sequence) and not isinstance( + sample_ids_value, (str, bytes) + ): + sample_ids = list(sample_ids_value) + else: + raise ValueError( + "CSD logical batch digest requires an ordered sample_id sequence." + ) + if len(sample_ids) != batch.size: + raise ValueError( + f"sample_id has {len(sample_ids)} entries for batch size {batch.size}." + ) + normalized_sample_ids = _canonical_logical_digest_json_value( + sample_ids, field="sample_id" + ) + if any( + sample_id is None or not str(sample_id).strip() + for sample_id in normalized_sample_ids + ): + raise ValueError("CSD logical batch digest found an empty sample_id.") + + field_names = [key for key in _CSD_LOGICAL_BATCH_BASE_FIELDS if key in batch] + field_names.extend( + sorted( + key + for key in batch + if key.startswith("teacher_") + and key.endswith(_CSD_LOGICAL_BATCH_TEACHER_SUFFIXES) + ) + ) + field_names.extend(sorted(key for key in batch if key.startswith("alignment_"))) + if not {"input_ids", "input_lengths", "token_mask", "sample_mask"}.issubset( + field_names + ): + raise ValueError( + "CSD logical batch digest is missing a required student field: " + f"available={tuple(field_names)}." + ) + if len(field_names) != len(set(field_names)): + raise AssertionError(f"Duplicate CSD logical digest fields: {field_names}.") + + sample_id_record = _logical_digest_field_record( + normalized_sample_ids, field="sample_id" + ) + batch_item_ids = tuple((batch_uid << 32) | ordinal for ordinal in range(batch.size)) + record: dict[str, Any] = { + "schema_version": _CSD_LOGICAL_BATCH_DIGEST_SCHEMA_VERSION, + "batch_uid": batch_uid, + "logical_samples": batch.size, + "sample_ids": sample_id_record, + "batch_item_ids": _logical_digest_field_record( + batch_item_ids, field="batch_item_id" + ), + "fields": { + key: _logical_digest_field_record(batch[key], field=key) + for key in field_names + }, + } + canonical_record = json.dumps( + record, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + record["record_sha256"] = hashlib.sha256(canonical_record).hexdigest() + return record + + +def log_csd_logical_batch_digest( + batch: BatchedDataDict[Any], *, batch_uid: int +) -> dict[str, Any]: + """Emit and return one compact arm-independent logical-batch record.""" + record = build_csd_logical_batch_digest_record(batch, batch_uid=batch_uid) + print( + "CSD_LOGICAL_BATCH_DIGEST " + + json.dumps( + record, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ), + flush=True, + ) + return record + + # =============================================================================== # Configuration # =============================================================================== @@ -133,6 +456,10 @@ class OffPolicyDistillationConfig(TypedDict): val_period: Validation cadence in steps. ``0`` disables validation. val_at_start: Run validation before training begins. val_at_end: Run validation on the final step. + offload_student_after_step: Offload the student model and optimizer to + CPU after each optimizer step. Student and teachers are colocated on + the same GPUs, so the next step's teacher forward otherwise competes + with resident student optimizer state. """ num_prompts_per_step: int @@ -142,6 +469,7 @@ class OffPolicyDistillationConfig(TypedDict): val_period: int val_at_start: bool val_at_end: bool + offload_student_after_step: NotRequired[bool] class OffPolicyDistillationSaveState(TypedDict): @@ -150,6 +478,7 @@ class OffPolicyDistillationSaveState(TypedDict): total_steps: int consumed_samples: int total_valid_tokens: int + next_packing_batch_uid: NotRequired[int] val_loss: NotRequired[float] @@ -160,6 +489,7 @@ def _default_off_policy_distillation_save_state() -> OffPolicyDistillationSaveSt "total_steps": 0, "consumed_samples": 0, "total_valid_tokens": 0, + "next_packing_batch_uid": 0, } @@ -173,10 +503,224 @@ class TeacherAlignerConfig(BaseModel, extra="allow"): drop_first_assistant_chunk_kl: Whether chat-mode alignment drops the first content pair in each assistant message for this teacher. The CE token mask is unaffected. + student_to_teacher_table_path: Path to this teacher's forward pseudo-target table + (student-to-teacher sub-token chains). ``None`` for a same-tokenizer + teacher. + teacher_to_student_table_path: Path to this teacher's reverse pseudo-target + table (teacher-to-student sub-token chains). ``None`` for a + same-tokenizer teacher. """ projection_matrix_path: Optional[str] = None drop_first_assistant_chunk_kl: bool = False + student_to_teacher_table_path: Optional[str] = None + teacher_to_student_table_path: Optional[str] = None + + +def _teacher_topk_ipc_keep_realized( + loss_config: CrossTokenizerDistillationLossConfig, +) -> bool: + """Whether sparse teacher IPC must retain realized teacher labels.""" + topk_k = int(loss_config["teacher_topk_ipc_k"]) + return topk_k > 0 and ( + bool(loss_config["csd_pure_alm"]) + or bool(loss_config["csd_realized_only_mismatch"]) + or topk_k == 1 + or bool(loss_config["teacher_topk_ipc_keep_realized"]) + ) + + +def _build_teacher_force_include_token_ids( + batch: BatchedDataDict[Any], + *, + teacher_idx: int, + loss_config: CrossTokenizerDistillationLossConfig, +) -> Optional[torch.Tensor]: + """Build realized teacher labels that sparse IPC must retain per position.""" + needs_force_ids = _teacher_topk_ipc_keep_realized(loss_config) or ( + int(loss_config["csd_noise_filter_topk"]) > 0 + ) + if not needs_force_ids: + return None + + teacher_input_ids = batch[f"teacher_{teacher_idx}_input_ids"] + if not torch.is_tensor(teacher_input_ids): + return None + + pair_valid = batch[f"alignment_{teacher_idx}_pair_valid"].to(torch.bool) + max_pairs = int(pair_valid.shape[1]) + student_spans = _chunk_ids_to_spans( + batch[f"alignment_{teacher_idx}_student_chunk_id"], max_pairs + ) + teacher_spans = _chunk_ids_to_spans( + batch[f"alignment_{teacher_idx}_teacher_chunk_id"], max_pairs + ) + force_ids = torch.full_like(teacher_input_ids, -1) + kl_chunk_shift = bool(loss_config["kl_chunk_shift"]) + batch_size, teacher_seq_len = teacher_input_ids.shape + + for batch_idx in range(batch_size): + for pair_idx in range(max_pairs): + if not bool(pair_valid[batch_idx, pair_idx]): + continue + student_start = int(student_spans[batch_idx, pair_idx, 0]) + student_end = int(student_spans[batch_idx, pair_idx, 1]) + teacher_start = int(teacher_spans[batch_idx, pair_idx, 0]) + teacher_end = int(teacher_spans[batch_idx, pair_idx, 1]) + if student_end <= student_start or teacher_end <= teacher_start: + continue + if kl_chunk_shift and student_start > 0 and teacher_start > 0: + predictor_start = teacher_start - 1 + else: + predictor_start = teacher_start + for offset in range(teacher_end - teacher_start): + predictor_pos = predictor_start + offset + label_pos = teacher_start + offset + if ( + 0 <= predictor_pos < teacher_seq_len + and 0 <= label_pos < teacher_seq_len + ): + force_ids[batch_idx, predictor_pos] = teacher_input_ids[ + batch_idx, label_pos + ] + return force_ids + + +# Ordered, non-overlapping alignment chunks request at most two distinct +# realized labels from one teacher row (see the builder below). +SPARSE_TEACHER_FORCE_WIDTH = 2 + + +def _build_teacher_force_include_token_id_slots( + batch: BatchedDataDict[Any], + *, + teacher_idx: int, + kl_chunk_shift: bool, +) -> torch.Tensor: + """Map every CSD teacher predictor row to the realized labels it must keep. + + Returns ``[B, T_t, SPARSE_TEACHER_FORCE_WIDTH]`` teacher label ids with + ``-1`` in unused slots. A chunk takes next-token predictors only when both + of its spans start after position zero (the CSD ``kl_chunk_shift`` rule), + so the last predictor of an unshifted chunk at the sequence start can also + be the first predictor of the following shifted chunk; those two requests + need different labels (row ``r`` is only ever asked for the token at ``r`` + or ``r + 1``). Chunks whose ``pair_valid`` or ``sample_mask`` is off, or + with an empty side (``[0, 0)`` span), request no label. + """ + teacher_input_ids = batch[f"teacher_{teacher_idx}_input_ids"] + pair_valid = batch[f"alignment_{teacher_idx}_pair_valid"].to(torch.bool) + pair_valid = pair_valid & batch["sample_mask"].to(torch.bool).unsqueeze(-1) + max_pairs = int(pair_valid.shape[1]) + student_spans = _chunk_ids_to_spans( + batch[f"alignment_{teacher_idx}_student_chunk_id"], max_pairs + ) + teacher_spans = _chunk_ids_to_spans( + batch[f"alignment_{teacher_idx}_teacher_chunk_id"], max_pairs + ) + batch_size, teacher_seq_len = teacher_input_ids.shape + force = torch.full( + (batch_size, teacher_seq_len, SPARSE_TEACHER_FORCE_WIDTH), + -1, + dtype=teacher_input_ids.dtype, + device=teacher_input_ids.device, + ) + # Driver-side and small: explicit loops keep the conflict errors precise. + for batch_idx, pair_idx in pair_valid.nonzero().tolist(): + student_start, student_end = student_spans[batch_idx, pair_idx].tolist() + teacher_start, teacher_end = teacher_spans[batch_idx, pair_idx].tolist() + if student_end <= student_start or teacher_end <= teacher_start: + continue + shift = 1 if kl_chunk_shift and student_start > 0 and teacher_start > 0 else 0 + for label_pos in range(teacher_start, teacher_end): + predictor_pos = label_pos - shift + label_id = int(teacher_input_ids[batch_idx, label_pos]) + row = force[batch_idx, predictor_pos] + if bool((row == label_id).any()): + continue + free_slots = (row == -1).nonzero() + if free_slots.numel() == 0: + raise ValueError( + "CSD alignment requests more than " + f"{SPARSE_TEACHER_FORCE_WIDTH} distinct realized labels from " + f"one teacher row; chunks must be ordered and non-overlapping. " + f"teacher={teacher_idx}, batch={batch_idx}, " + f"pos={predictor_pos}, existing={row.tolist()}, new={label_id}." + ) + row[int(free_slots[0, 0])] = label_id + return force + + +def _get_teacher_logits_ipc( + teacher_policy: Policy, + teacher_data: BatchedDataDict[Any], + loss_fn: CrossTokenizerDistillationLossFn, + *, + teacher_idx: int, + micro_batch_size: int, + force_include_token_ids: Optional[torch.Tensor], + timer: Optional[Timer], + packing_plan: Optional[LockstepPackingPlan] = None, + packing_side_id: Optional[str] = None, +) -> tuple[str, list[dict[str, Any]]]: + """Return the indexed transport suffix and IPC handles for one teacher.""" + topk_k = ( + int(loss_fn.cfg.teacher_topk_ipc_k) if isinstance(loss_fn, CSDLossFn) else 0 + ) + if packing_plan is not None and topk_k > 0: + raise ValueError( + "CSD lockstep packing currently supports dense teacher IPC only." + ) + if topk_k > 0 and _loss_teacher_is_cross_tokenizer(loss_fn, teacher_idx): + noise_filter_topk = int(loss_fn.cfg.csd_noise_filter_topk) + if noise_filter_topk < 0: + raise ValueError( + f"csd_noise_filter_topk must be >= 0, got {noise_filter_topk}" + ) + if force_include_token_ids is not None: + teacher_data["force_include_token_ids"] = force_include_token_ids + if loss_fn.cfg.owner_local: + # Megatron teachers export the sparse rows of each CP shard, with + # every realized label the owner-local loss requests. + if force_include_token_ids is None: + raise ValueError( + f"Teacher {teacher_idx} sparse rows need forced labels." + ) + handles = teacher_policy.get_sparse_logits_ipc( + teacher_data, + k=topk_k, + temperature=float(loss_fn.cfg.temperature), + real_vocab_size=int(loss_fn.teacher_vocab_sizes[teacher_idx]), + noise_filter_k=noise_filter_topk, + micro_batch_size=micro_batch_size, + timer=timer, + ) + return "sparse_logits_ipc", handles + handles = teacher_policy.get_topk_logits_ipc( + teacher_data, + k=topk_k, + temperature=float(loss_fn.cfg.temperature), + vocab_size=int(loss_fn.teacher_vocab_sizes[teacher_idx]), + micro_batch_size=micro_batch_size, + support_mode=str(loss_fn.cfg.teacher_topk_ipc_support_mode), + gt_filter_topk=(noise_filter_topk if noise_filter_topk > 0 else None), + timer=timer, + ) + return "sparse_logits_ipc", handles + + full_logits_kwargs: dict[str, Any] = {} + if packing_plan is not None: + full_logits_kwargs = { + "packing_plan": packing_plan, + "packing_side_id": packing_side_id, + } + handles = teacher_policy.get_full_logits_ipc( + teacher_data, + micro_batch_size=micro_batch_size, + timer=timer, + **full_logits_kwargs, + ) + return "full_logits_ipc", handles class TeacherConfig(BaseModel, extra="allow"): @@ -193,11 +737,8 @@ class TeacherConfig(BaseModel, extra="allow"): weight: Static loss weight for this teacher when several teachers are aggregated (``kd_loss_mode="sum"`` / the convex ``"averaged_logits"`` mix). Single-teacher runs leave it at ``1.0``. - gold_loss: Optional per-teacher override of ``loss_fn.gold_loss``. - ``None`` falls back to the global value. Honored only in - ``kd_loss_mode="sum"`` (other modes use the global). - xtoken_loss: Optional per-teacher override of ``loss_fn.xtoken_loss``, - same semantics as ``gold_loss``. + gold_loss: Optional X-Token per-teacher override of loss_fn.gold_loss. + xtoken_loss: Optional X-Token per-teacher override of loss_fn.xtoken_loss. """ aligner: TeacherAlignerConfig = Field(default_factory=TeacherAlignerConfig) @@ -205,6 +746,20 @@ class TeacherConfig(BaseModel, extra="allow"): gold_loss: Optional[bool] = None xtoken_loss: Optional[bool] = None + @model_validator(mode="before") + @classmethod + def _reject_misplaced_csd_tables(cls, value: Any) -> Any: + if isinstance(value, dict): + for key in ( + "student_to_teacher_table_path", + "teacher_to_student_table_path", + ): + if key in value: + raise ValueError( + f"teachers[i].{key} must be nested under teachers[i].aligner.{key}" + ) + return value + def policy_config(self) -> PolicyConfig: """Recover the plain ``PolicyConfig`` dict (cross-tokenizer knobs stripped).""" return cast( @@ -214,6 +769,7 @@ def policy_config(self) -> PolicyConfig: class MasterConfig(BaseModel, extra="allow"): + algorithm: Literal["xtoken"] = "xtoken" policy: PolicyConfig # student teachers: list[TeacherConfig] = Field(min_length=1) loss_fn: CrossTokenizerDistillationLossConfig @@ -227,11 +783,695 @@ class MasterConfig(BaseModel, extra="allow"): checkpointing: CheckpointingConfig +class CSDMasterConfig(MasterConfig): + """CSD configuration using the shared cross-tokenizer training workflow.""" + + algorithm: Literal["csd"] = "csd" + loss_fn: CSDLossConfig + + @model_validator(mode="after") + def _validate_csd_tables(self) -> "CSDMasterConfig": + for i, teacher in enumerate(self.teachers): + aligner = teacher.aligner + forward = aligner.student_to_teacher_table_path + reverse = aligner.teacher_to_student_table_path + if bool(forward) != bool(reverse): + raise ValueError( + f"teachers[{i}].aligner requires both student_to_teacher_table_path " + "and teacher_to_student_table_path for CSD." + ) + if aligner.projection_matrix_path and not forward: + raise ValueError(f"teachers[{i}]: CSD requires both sub-token tables.") + if ( + forward + and not self.loss_fn.common_indices_from_subtoks + and not aligner.projection_matrix_path + ): + raise ValueError( + f"teachers[{i}].aligner.projection_matrix_path is required " + "when loss_fn.common_indices_from_subtoks=false; set " + "common_indices_from_subtoks=true to use the forward table " + "for common-token mapping." + ) + return self + + +def _teacher_is_cross_tokenizer(config: MasterConfig, teacher: TeacherConfig) -> bool: + """Select alignment using the artifacts belonging to the chosen algorithm.""" + if config.algorithm == "csd": + return teacher.aligner.student_to_teacher_table_path is not None + return teacher.aligner.projection_matrix_path is not None + + +def _loss_teacher_is_cross_tokenizer( + loss_fn: CrossTokenizerDistillationLossFn | CSDLossFn, i: int +) -> bool: + if isinstance(loss_fn, CSDLossFn): + return not loss_fn._teacher_is_same_vocab(i) + return loss_fn.projection_matrix_paths[i] is not None + + +def _loss_config_dict( + config: MasterConfig, +) -> CrossTokenizerDistillationLossConfig | dict[str, Any]: + if isinstance(config.loss_fn, CSDLossConfig): + return config.loss_fn.model_dump() + return config.loss_fn + + # =============================================================================== # Setup # =============================================================================== +def _xtoken_entity_parallelism( + cfg: PolicyConfig, *, label: str +) -> tuple[int, int, int]: + """Return ``(tensor_parallel, context_parallel, pipeline_parallel)`` for a config. + + Cross-tokenizer execution uses either the Automodel backend + (``automodel_cfg.enabled``) or the Megatron backend + (``megatron_cfg.enabled``). Both support tensor and context parallelism (the + loss is parallelism-invariant); Megatron additionally supports pipeline + parallelism, where only the last stage holds logits and therefore only it + contributes teacher full-logits IPC handles. DTensor has no pipeline axis, + so it always reports ``pipeline_parallel == 1``. + + Args: + cfg: The policy or teacher policy config. + label: Human-readable name of the entity for error messages, e.g. + ``"teachers[0]"``. + + Returns: + The entity's tensor-, context- and pipeline-parallel sizes, read from + whichever backend block is enabled. + """ + if "megatron_cfg" in cfg and cfg["megatron_cfg"]["enabled"]: + mcfg = cfg["megatron_cfg"] + return ( + mcfg["tensor_model_parallel_size"], + mcfg["context_parallel_size"], + mcfg["pipeline_model_parallel_size"], + ) + assert cfg["automodel_cfg"]["enabled"], ( + f"Cross-tokenizer distillation requires {label} to enable either Automodel " + f"(automodel_cfg.enabled=true) or Megatron " + f"(megatron_cfg.enabled=true)." + ) + return ( + cfg["automodel_cfg"]["tensor_parallel_size"], + cfg["automodel_cfg"]["context_parallel_size"], + 1, + ) + + +def validate_csd_router_replay_setup(master_config: MasterConfig) -> None: + """Fail before cluster creation when CSD replay cannot be consumed.""" + if not router_replay_enabled(master_config.policy): + return + megatron_cfg = master_config.policy.get("megatron_cfg") or {} + if not megatron_cfg.get("enabled", False): + raise ValueError( + "policy.router_replay.enabled=true for CSD distillation requires " + "the Megatron student policy backend; DTensor cannot consume " + "routed_experts." + ) + + +def _sequence_packing_enabled(cfg: PolicyConfig) -> bool: + packing = cfg.get("sequence_packing") or {} + return bool(packing.get("enabled", False)) + + +def _megatron_enabled(cfg: PolicyConfig) -> bool: + return "megatron_cfg" in cfg and bool(cfg["megatron_cfg"]["enabled"]) + + +def csd_owner_local_ineligibility( + master_config: MasterConfig, loss_config: CSDLossConfig +) -> list[str]: + """Reasons a CSD run cannot evaluate cross-tokenizer KD owner-local. + + Owner-local evaluation keeps the student logits in their context-parallel + shard and reads only the teacher rows each rank's chunks need. It needs a + Megatron student with PP=1 and neither sequence packing nor dynamic + batching, loss settings listed by :func:`owner_local_config_violations`, + and, with sparse teacher IPC, Megatron cross-tokenizer teachers whose + sequence width matches their CP shards. An empty list means eligible. + """ + policy_config = master_config.policy + if not _megatron_enabled(policy_config): + return ["the student policy is not Megatron"] + reasons: list[str] = [] + _, _, student_pp = _xtoken_entity_parallelism(policy_config, label="policy") + if student_pp != 1: + reasons.append(f"policy pipeline_model_parallel_size={student_pp} (needs 1)") + if _sequence_packing_enabled(policy_config): + reasons.append("policy sequence packing is enabled") + dynamic_batching = policy_config.get("dynamic_batching") or {} + if bool(dynamic_batching.get("enabled", False)): + reasons.append("policy dynamic batching is enabled") + reasons.extend(owner_local_config_violations(loss_config)) + if loss_config.teacher_topk_ipc_k > 0: + for i, teacher in enumerate(master_config.teachers): + if not _teacher_is_cross_tokenizer(master_config, teacher): + continue + teacher_config = teacher.policy_config() + if not _megatron_enabled(teacher_config): + reasons.append( + f"teachers[{i}] is not Megatron (sparse teacher rows come " + "from the Megatron sender)" + ) + continue + tp, cp, _ = _xtoken_entity_parallelism( + teacher_config, label=f"teachers[{i}]" + ) + divisor = 2 * cp if cp > 1 else 1 + if tp > 1 and bool(teacher_config["megatron_cfg"]["sequence_parallel"]): + divisor = math.lcm(divisor, tp) + if int(teacher_config["make_sequence_length_divisible_by"]) % divisor: + reasons.append( + f"teachers[{i}].make_sequence_length_divisible_by=" + f"{teacher_config['make_sequence_length_divisible_by']} is not " + f"a multiple of {divisor} (2 x CP, and TP with sequence " + "parallelism)" + ) + return reasons + + +def _positive_config_int(value: object, *, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{field} must be a positive integer, got {value!r}.") + return value + + +def _csd_configured_dp_size( + cfg: PolicyConfig, + *, + label: str, + world_size: int, +) -> int: + tp, cp, pp = _xtoken_entity_parallelism(cfg, label=label) + model_parallel_size = tp * cp * pp + if world_size % model_parallel_size != 0: + raise ValueError( + f"CSD lockstep packing cannot form {label}'s DP grid: world_size=" + f"{world_size} is not divisible by tp*cp*pp={model_parallel_size}." + ) + return world_size // model_parallel_size + + +def _load_csd_packed_model_config(cfg: PolicyConfig) -> Any: + """Resolve architecture metadata without constructing a worker or model. + + A local tokenizer snapshot is preferred when the model is still named by a + mutable Hub id and that snapshot also contains ``config.json``. This keeps + qualification architecture checks on the same pinned revision used by the + collator and backend tokenizer. Local model/checkpoint paths remain + authoritative. Any format whose architecture cannot be resolved is denied + below rather than silently admitted to recurrent-state packing. + """ + from transformers import AutoConfig + + model_source = str(cfg.get("model_name", "")) + tokenizer_source = str((cfg.get("tokenizer") or {}).get("name", "")) + architecture_source = model_source + if ( + not os.path.exists(model_source) + and os.path.isdir(tokenizer_source) + and os.path.isfile(os.path.join(tokenizer_source, "config.json")) + ): + architecture_source = tokenizer_source + + hf_config_overrides = cfg.get("hf_config_overrides") or {} + return AutoConfig.from_pretrained( + architecture_source, + trust_remote_code=True, + local_files_only=os.getenv("HF_HUB_OFFLINE") == "1", + **hf_config_overrides, + ) + + +def _iter_nonempty_recurrent_config_fields( + value: object, *, prefix: str = "" +) -> Iterator[str]: + """Yield configured Mamba/SSM field paths from nested model metadata.""" + if not isinstance(value, Mapping): + return + for key, child in value.items(): + field = f"{prefix}.{key}" if prefix else str(key) + key_lower = str(key).lower() + if ("mamba" in key_lower or "ssm" in key_lower) and child not in ( + None, + False, + 0, + "", + [], + {}, + ): + yield field + yield from _iter_nonempty_recurrent_config_fields(child, prefix=field) + + +def _validate_csd_packed_model_architecture(cfg: PolicyConfig, *, label: str) -> None: + """Reject recurrent/hybrid architectures using resolved model metadata.""" + model_name = str(cfg.get("model_name", "")) + try: + model_config = _load_csd_packed_model_config(cfg) + except Exception as error: + raise ValueError( + "CSD lockstep packing must resolve model architecture metadata " + f"before cluster creation; could not inspect {label}.model_name=" + f"{model_name!r}: {type(error).__name__}: {error}" + ) from error + + metadata = model_config.to_dict() + identity_metadata = { + "model_type": getattr(model_config, "model_type", None), + "architectures": getattr(model_config, "architectures", None), + "layer_types": metadata.get("layer_types"), + "layers_block_type": metadata.get("layers_block_type"), + "hybrid_override_pattern": metadata.get("hybrid_override_pattern"), + } + identity_text = repr(identity_metadata).lower() + unsafe_marker = next( + (marker for marker in _UNSAFE_PACKED_MODEL_MARKERS if marker in identity_text), + None, + ) + recurrent_fields = tuple(_iter_nonempty_recurrent_config_fields(metadata)) + hybrid_pattern = metadata.get("hybrid_override_pattern") + if unsafe_marker is not None or recurrent_fields or hybrid_pattern: + raise ValueError( + "CSD lockstep packing does not yet support Nano/Mamba recurrent " + f"state isolation for {label}.model_name={model_name!r}; resolved " + f"architecture={identity_metadata!r}, recurrent_fields=" + f"{recurrent_fields!r}." + ) + + +def _validate_csd_packed_backend(cfg: PolicyConfig, *, label: str) -> None: + """Validate the initial dense-transformer packing support matrix.""" + model_name = str(cfg.get("model_name", "")).lower() + unsafe_marker = next( + (marker for marker in _UNSAFE_PACKED_MODEL_MARKERS if marker in model_name), + None, + ) + if unsafe_marker is not None: + raise ValueError( + f"CSD lockstep packing does not yet support Nano/Mamba state " + f"isolation; {label}.model_name={cfg.get('model_name')!r} matched " + f"{unsafe_marker!r}." + ) + _validate_csd_packed_model_architecture(cfg, label=label) + + make_divisible = _positive_config_int( + cfg.get("make_sequence_length_divisible_by"), + field=f"{label}.make_sequence_length_divisible_by", + ) + tp, cp, pp = _xtoken_entity_parallelism(cfg, label=label) + megatron_cfg = cfg.get("megatron_cfg") or {} + if bool(megatron_cfg.get("enabled", False)): + if pp != 1: + raise ValueError( + "CSD lockstep packing currently requires Megatron PP=1; " + f"{label} has pipeline_model_parallel_size={pp}." + ) + if cp not in (1, 2): + raise ValueError( + "CSD lockstep packing currently supports Megatron CP=1 or " + f"CP=2; {label} has context_parallel_size={cp}." + ) + fp8_cfg = megatron_cfg.get("fp8_cfg") or {} + if bool(fp8_cfg.get("enabled", False)): + raise ValueError( + "CSD lockstep packing has not qualified Megatron FP8 packed-" + f"total padding; disable {label}.megatron_cfg.fp8_cfg.enabled." + ) + if ( + megatron_cfg.get("moe_token_dispatcher_type") == "flex" + and megatron_cfg.get("moe_flex_dispatcher_backend") == "hybridep" + ): + raise ValueError( + "CSD lockstep packing has not qualified HybridEP packed-total " + f"padding for {label}." + ) + minimum_divisor = cp * 2 if cp > 1 else 1 + if tp > 1 and bool(megatron_cfg.get("sequence_parallel", False)): + minimum_divisor *= tp + if make_divisible % minimum_divisor != 0: + raise ValueError( + f"{label}.make_sequence_length_divisible_by={make_divisible} " + f"must be a multiple of {minimum_divisor} for Megatron " + f"TP={tp}, CP={cp}, sequence_parallel=" + f"{bool(megatron_cfg.get('sequence_parallel', False))}." + ) + return + + automodel_cfg = cfg["automodel_cfg"] + if cp != 1: + raise ValueError( + "CSD lockstep packing currently supports Automodel CP=1 only; " + f"{label} has context_parallel_size={cp}." + ) + if bool(automodel_cfg.get("sequence_parallel", False)): + raise ValueError( + "CSD lockstep packing currently requires Automodel sequence " + f"parallelism to be disabled for {label}." + ) + + +def validate_csd_packing_setup(master_config: MasterConfig) -> Optional[int]: + """Fail before cluster creation for unsupported CSD packing configs. + + Returns the common configured DP size when packing is enabled and ``None`` + when every side has packing disabled. + """ + policy_config = master_config.policy + teacher_configs = [teacher.policy_config() for teacher in master_config.teachers] + side_configs: list[tuple[str, PolicyConfig]] = [ + ("policy", policy_config), + *[ + (f"teachers[{teacher_idx}]", teacher_config) + for teacher_idx, teacher_config in enumerate(teacher_configs) + ], + ] + packing_modes = [_sequence_packing_enabled(config) for _, config in side_configs] + if not any(packing_modes): + return None + if master_config.algorithm != "csd": + raise ValueError("Lockstep sequence packing requires algorithm=csd.") + if not all(packing_modes): + enabled_labels = [ + label + for (label, _), enabled in zip(side_configs, packing_modes, strict=True) + if enabled + ] + disabled_labels = [ + label + for (label, _), enabled in zip(side_configs, packing_modes, strict=True) + if not enabled + ] + raise ValueError( + "CSD lockstep packing must be enabled on the student and every " + f"teacher; enabled={enabled_labels}, disabled={disabled_labels}." + ) + + data_config = master_config.data + if master_config.collator.num_packed_rows != 1: + raise ValueError( + "CSD lockstep packing keeps one source row per logical sample; " + f"collator.num_packed_rows must be 1, got " + f"{master_config.collator.num_packed_rows!r}." + ) + default_config = data_config.get("default") or {} + for split_name in ("train", "validation"): + split_configs = data_config.get(split_name) + if split_configs is None: + continue + if isinstance(split_configs, Mapping): + dataset_configs = [split_configs] + elif isinstance(split_configs, list): + dataset_configs = [ + cast(Mapping[str, Any], dataset_config) + for dataset_config in split_configs + ] + else: + raise ValueError( + f"data.{split_name} must be a dataset config or list of configs, " + f"got {type(split_configs).__name__}." + ) + for dataset_idx, dataset_config in enumerate(dataset_configs): + characters_per_sample = dataset_config.get( + "characters_per_sample", default_config.get("characters_per_sample") + ) + if characters_per_sample is not None: + raise ValueError( + "CSD lockstep packing requires one Arrow/chat row per " + "logical sample; set " + f"data.{split_name}.characters_per_sample=null " + f"(dataset {dataset_idx} resolved to " + f"{characters_per_sample!r})." + ) + + algorithms: list[str] = [] + for label, config in side_configs: + dynamic_batching = config.get("dynamic_batching") or {} + if bool(dynamic_batching.get("enabled", False)): + raise ValueError( + "CSD lockstep packing is incompatible with independent " + f"dynamic batching; disable {label}.dynamic_batching.enabled." + ) + sequence_packing = config.get("sequence_packing") or {} + algorithm = str(sequence_packing.get("algorithm", "")) + algorithms.append(algorithm) + if "fuse_loss" not in sequence_packing: + raise ValueError( + f"{label}.sequence_packing.fuse_loss must be set explicitly; " + "CSD lockstep packing requires false." + ) + if bool(sequence_packing["fuse_loss"]): + raise ValueError( + "CSD lockstep packing requires sequence_packing.fuse_loss=false; " + f"it is enabled for {label}." + ) + _positive_config_int( + sequence_packing.get("train_mb_tokens"), + field=f"{label}.sequence_packing.train_mb_tokens", + ) + if label.startswith("teachers["): + _positive_config_int( + sequence_packing.get("logprob_mb_tokens"), + field=f"{label}.sequence_packing.logprob_mb_tokens", + ) + if config.get("train_micro_batch_size") != 1: + raise ValueError( + "CSD lockstep packing preserves MBS=1 logical loss cohorts; " + f"{label}.train_micro_batch_size must be 1, got " + f"{config.get('train_micro_batch_size')!r}." + ) + _validate_csd_packed_backend(config, label=label) + + if set(algorithms) != {_CSD_LOCKSTEP_PACKING_ALGORITHM}: + raise ValueError( + "CSD lockstep packing selects one controller algorithm for every " + f"side; expected {_CSD_LOCKSTEP_PACKING_ALGORITHM!r}, got " + f"{dict(zip((label for label, _ in side_configs), algorithms, strict=True))}." + ) + + loss_config = _loss_config_dict(master_config) + if int(loss_config.get("teacher_topk_ipc_k", 0)) != 0: + raise ValueError( + "CSD lockstep packing currently supports dense teacher IPC only; " + "set loss_fn.teacher_topk_ipc_k=0." + ) + if loss_config.get("kd_loss_mode") != "sum": + raise ValueError( + "CSD lockstep packing currently supports static additive teacher " + "aggregation only; set loss_fn.kd_loss_mode='sum'." + ) + if loss_config.get("sum_weights_metric") is not None: + raise ValueError( + "CSD lockstep packing does not support per-batch teacher scoring; " + "set loss_fn.sum_weights_metric=null." + ) + + world_size = _positive_config_int( + master_config.cluster.num_nodes, field="cluster.num_nodes" + ) * _positive_config_int( + master_config.cluster.gpus_per_node, field="cluster.gpus_per_node" + ) + dp_sizes = [ + _csd_configured_dp_size(config, label=label, world_size=world_size) + for label, config in side_configs + ] + if len(set(dp_sizes)) != 1: + raise ValueError( + "CSD lockstep packing initially requires equal student/teacher " + f"DP sizes; got {dict(zip((label for label, _ in side_configs), dp_sizes, strict=True))}." + ) + data_parallel_size = dp_sizes[0] + global_batch_size = _positive_config_int( + master_config.distillation["num_prompts_per_step"], + field="distillation.num_prompts_per_step", + ) + if global_batch_size % data_parallel_size != 0: + raise ValueError( + f"CSD logical GBS={global_batch_size} must be divisible by the " + f"common DP size {data_parallel_size}." + ) + for label, config in side_configs: + if config.get("train_global_batch_size") != global_batch_size: + raise ValueError( + "CSD GBS counts logical samples and must match on every side; " + f"distillation GBS={global_batch_size}, " + f"{label}.train_global_batch_size=" + f"{config.get('train_global_batch_size')!r}." + ) + + student_tp, student_cp, student_pp = _xtoken_entity_parallelism( + policy_config, label="policy" + ) + for teacher_idx, teacher_config in enumerate(teacher_configs): + teacher_tp, teacher_cp, teacher_pp = _xtoken_entity_parallelism( + teacher_config, label=f"teachers[{teacher_idx}]" + ) + assert_xtoken_ipc_node_local( + num_nodes=master_config.cluster.num_nodes, + gpus_per_node=master_config.cluster.gpus_per_node, + student_tp=student_tp, + student_cp=student_cp, + student_pp=student_pp, + teacher_tp=teacher_tp, + teacher_cp=teacher_cp, + teacher_pp=teacher_pp, + student_dp=data_parallel_size, + teacher_dp=data_parallel_size, + ) + return data_parallel_size + + +def _freeze_tokenizer_value(value: object) -> object: + """Convert tokenizer metadata into a deterministic comparison value.""" + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, Mapping): + return tuple( + sorted( + (str(key), _freeze_tokenizer_value(item)) for key, item in value.items() + ) + ) + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return tuple(_freeze_tokenizer_value(item) for item in value) + added_token_fields = ( + "content", + "single_word", + "lstrip", + "rstrip", + "normalized", + "special", + ) + if all(hasattr(value, field) for field in added_token_fields): + return tuple( + (field, _freeze_tokenizer_value(getattr(value, field))) + for field in added_token_fields + ) + return repr(value) + + +def _tokenizer_reuse_signature( + tokenizer: PreTrainedTokenizerBase, + tokenizer_config: Mapping[str, Any], +) -> dict[str, object]: + """Describe every tokenizer semantic required for safe token-ID reuse.""" + try: + vocab = tokenizer.get_vocab() + except Exception as error: + raise ValueError( + "same-tokenizer reuse requires tokenizer.get_vocab() to be available" + ) from error + if not isinstance(vocab, Mapping): + raise ValueError( + "same-tokenizer reuse requires tokenizer.get_vocab() to return a mapping" + ) + try: + vocab_items = tuple( + sorted((str(token), int(idx)) for token, idx in vocab.items()) + ) + except (TypeError, ValueError) as error: + raise ValueError( + "same-tokenizer reuse requires a string-to-integer vocabulary mapping" + ) from error + + backend_signature: object = None + is_fast = bool(getattr(tokenizer, "is_fast", False)) + if not is_fast: + raise ValueError( + "same-tokenizer reuse requires fast tokenizers with a complete " + "serialized backend signature; slow-tokenizer normalization, " + "pretokenization, and postprocessing cannot be proven equivalent" + ) + backend_tokenizer = getattr(tokenizer, "backend_tokenizer", None) + if backend_tokenizer is None: + raise ValueError( + "same-tokenizer reuse requires a serializable fast-tokenizer backend" + ) + to_str = getattr(backend_tokenizer, "to_str", None) + if not callable(to_str): + raise ValueError( + "same-tokenizer reuse requires fast tokenizer backend serialization" + ) + backend_signature = to_str() + + return { + "tokenizer_class": ( + type(tokenizer).__module__, + type(tokenizer).__qualname__, + ), + "is_fast": is_fast, + "vocab": vocab_items, + "backend": backend_signature, + "special_tokens_map": _freeze_tokenizer_value( + getattr(tokenizer, "special_tokens_map_extended", {}) + ), + "all_special_ids": tuple(int(idx) for idx in tokenizer.all_special_ids), + "chat_template": getattr(tokenizer, "chat_template", None), + "chat_template_kwargs": _freeze_tokenizer_value( + tokenizer_config.get("chat_template_kwargs") or {} + ), + } + + +def _assert_same_tokenizer_reuse_safe( + *, + student_tokenizer: PreTrainedTokenizerBase, + teacher_tokenizer: PreTrainedTokenizerBase, + student_config: PolicyConfig, + teacher_config: PolicyConfig, + teacher_idx: int, +) -> None: + student_signature = _tokenizer_reuse_signature( + student_tokenizer, student_config["tokenizer"] + ) + teacher_signature = _tokenizer_reuse_signature( + teacher_tokenizer, teacher_config["tokenizer"] + ) + differing_fields = sorted( + field + for field in student_signature + if student_signature[field] != teacher_signature[field] + ) + if differing_fields: + raise ValueError( + f"teachers[{teacher_idx}].aligner.projection_matrix_path is null, " + "but safe " + "student-token reuse requires identical tokenizer mapping, backend, " + "special-token configuration, chat template, and template kwargs; " + f"differing fields: {differing_fields}." + ) + + +def validate_csd_tokenizer_reuse( + master_config: MasterConfig, + student_tokenizer: PreTrainedTokenizerBase, + teacher_tokenizers: Sequence[PreTrainedTokenizerBase], +) -> None: + """Validate every null-projection teacher before any worker is created.""" + if len(teacher_tokenizers) != len(master_config.teachers): + raise ValueError( + f"expected one tokenizer per teacher; got {len(teacher_tokenizers)} " + f"tokenizers for {len(master_config.teachers)} teachers." + ) + for teacher_idx, teacher in enumerate(master_config.teachers): + if not _teacher_is_cross_tokenizer(master_config, teacher): + _assert_same_tokenizer_reuse_safe( + student_tokenizer=student_tokenizer, + teacher_tokenizer=teacher_tokenizers[teacher_idx], + student_config=master_config.policy, + teacher_config=teacher.policy_config(), + teacher_idx=teacher_idx, + ) + + def setup( master_config: MasterConfig, student_tokenizer: PreTrainedTokenizerBase, @@ -253,7 +1493,7 @@ def setup( policy_config = master_config.policy teachers = master_config.teachers teacher_configs = [t.policy_config() for t in teachers] - loss_config = master_config.loss_fn + loss_config = _loss_config_dict(master_config) distillation_config = master_config.distillation data_config = master_config.data logger_config = master_config.logger @@ -264,46 +1504,141 @@ def setup( f"tokenizers for {len(teachers)} teachers." ) - # Backend gate: DTensor only, for the student and every teacher. Unlike - # the TP=CP=1 multi-teacher prototype, this path supports TP/CP/diff-DP - # sharding (the loss is parallelism-invariant), so there is deliberately NO - # tensor/context_parallel_size==1 assert. - assert (policy_config.get("automodel_cfg") or {}).get("enabled"), ( - "xtoken distillation requires policy.automodel_cfg.enabled=true." - ) + if master_config.algorithm == "xtoken": + for label, cfg in [ + ("policy", policy_config), + *[(f"teachers[{i}]", tc) for i, tc in enumerate(teacher_configs)], + ]: + if not cfg["automodel_cfg"]["enabled"]: + raise ValueError( + f"X-Token requires {label}.automodel_cfg.enabled=true." + ) + validate_csd_router_replay_setup(master_config) + + # Backend gate. The student and each teacher may be Automodel or Megatron. + # Both support TP/CP/diff-DP sharding and Megatron additionally supports PP + # (the loss is parallelism-invariant), so there is deliberately NO + # tensor/context/pipeline_parallel_size==1 assert on either backend. + # ``_xtoken_entity_parallelism`` validates each entity's backend as a side + # effect (raising on an unsupported / misconfigured one). + _xtoken_entity_parallelism(policy_config, label="policy") for i, tc in enumerate(teacher_configs): - assert (tc.get("automodel_cfg") or {}).get("enabled"), ( - f"xtoken distillation requires teachers.{i}.automodel_cfg.enabled=true." - ) - - # A null projection path marks a same-vocab teacher (direct KL, no - # projection/alignment); that only makes sense when it shares the student's - # vocab, so check vocab SIZE (tokenizer names differ for same-vocab pairs - # like Llama-3.2-3B vs -1B). Cross-tokenizer teachers (non-null path) are - # validated against their projection matrix in the loss. - student_vocab = len(student_tokenizer) - for i, teacher in enumerate(teachers): - if teacher.aligner.projection_matrix_path is None: - assert len(teacher_tokenizers[i]) == student_vocab, ( - f"teachers[{i}].aligner.projection_matrix_path is null " - f"(same-vocab teacher) but its tokenizer vocab " - f"({len(teacher_tokenizers[i])}) != student vocab " - f"({student_vocab}). A same-vocab teacher must share the " - "student tokenizer; set aligner.projection_matrix_path to run " - "it as a cross-tokenizer teacher." + _xtoken_entity_parallelism(tc, label=f"teachers[{i}]") + configured_packing_dp = validate_csd_packing_setup(master_config) + + teacher_topk_ipc_k = ( + int(loss_config["teacher_topk_ipc_k"]) + if master_config.algorithm == "csd" + else 0 + ) + if teacher_topk_ipc_k < 0: + raise ValueError( + f"loss_fn.teacher_topk_ipc_k must be >= 0, got {teacher_topk_ipc_k}" + ) + teacher_topk_support_mode = ( + str(loss_config["teacher_topk_ipc_support_mode"]) + if master_config.algorithm == "csd" + else "row_topk" + ) + if teacher_topk_support_mode != "row_topk": + raise ValueError( + "loss_fn.teacher_topk_ipc_support_mode currently supports only " + f"'row_topk', got {teacher_topk_support_mode!r}" + ) + if teacher_topk_ipc_k > 0: + has_sparse_teacher = any( + _teacher_is_cross_tokenizer(master_config, teacher) for teacher in teachers + ) + if has_sparse_teacher and loss_config["sum_weights_metric"] is not None: + raise ValueError( + "Sparse teacher IPC cannot be combined with sum_weights_metric; " + "dynamic teacher scoring requires full logits." ) + if has_sparse_teacher and loss_config["kd_loss_mode"] == "select_teacher": + raise ValueError( + "Sparse teacher IPC cannot be combined with " + "kd_loss_mode='select_teacher'; teacher selection requires full logits." + ) + if not _megatron_enabled(policy_config): + # The replicated loss reads the Automodel sparse format. + for i, (teacher, teacher_config) in enumerate( + zip(teachers, teacher_configs, strict=True) + ): + if not _teacher_is_cross_tokenizer(master_config, teacher): + continue + if not teacher_config["automodel_cfg"]["enabled"]: + raise ValueError( + "Sparse teacher IPC with an Automodel student requires " + "Automodel cross-tokenizer teachers; teachers[{i}] is not " + "Automodel. Set loss_fn.teacher_topk_ipc_k=0 for dense " + "teacher IPC.".format(i=i) + ) - # With context parallelism every forwarded sequence width must split - # evenly across the CP ranks, and batch widths follow the data. + # Megatron students evaluate cross-tokenizer KD on the context-parallel + # rank that owns each chunk whenever the run allows it; otherwise the loss + # stays replicated, which only dense teacher IPC can feed. + owner_local = False + if master_config.algorithm == "csd": + owner_local_reasons = csd_owner_local_ineligibility( + master_config, CSDLossConfig.model_validate(loss_config) + ) + if _megatron_enabled(policy_config): + if teacher_topk_ipc_k > 0 and owner_local_reasons: + raise ValueError( + "Sparse teacher IPC (loss_fn.teacher_topk_ipc_k > 0) with a " + "Megatron student needs owner-local CSD evaluation, which " + "this run cannot use: " + + "; ".join(owner_local_reasons) + + ". Set loss_fn.teacher_topk_ipc_k=0 for dense teacher IPC." + ) + owner_local = not owner_local_reasons + print( + "CSD loss evaluation: " + + ( + "owner-local" + if owner_local + else "replicated (" + "; ".join(owner_local_reasons) + ")" + ), + flush=True, + ) + + if master_config.algorithm == "csd": + validate_csd_tokenizer_reuse( + master_config, student_tokenizer, teacher_tokenizers + ) + else: + # A null projection path marks a same-vocab teacher (direct KL, no + # projection/alignment); that only makes sense when it shares the student's + # vocab, so check vocab SIZE (tokenizer names differ for same-vocab pairs + # like Llama-3.2-3B vs -1B). Cross-tokenizer teachers (non-null path) are + # validated against their projection matrix in the loss. + student_vocab = len(student_tokenizer) + for i, teacher in enumerate(teachers): + if teacher.aligner.projection_matrix_path is None: + assert len(teacher_tokenizers[i]) == student_vocab, ( + f"teachers[{i}].aligner.projection_matrix_path is null " + f"(same-vocab teacher) but its tokenizer vocab " + f"({len(teacher_tokenizers[i])}) != student vocab " + f"({student_vocab}). A same-vocab teacher must share the " + "student tokenizer; set aligner.projection_matrix_path to run " + "it as a cross-tokenizer teacher." + ) + + # With Automodel context parallelism every forwarded sequence width must + # split evenly across the CP ranks, and batch widths follow the data. + # Megatron pads its own inputs to its context-parallel layout. student_div = policy_config["make_sequence_length_divisible_by"] - student_cp = policy_config["automodel_cfg"]["context_parallel_size"] - assert student_div % student_cp == 0, ( - f"policy.make_sequence_length_divisible_by ({student_div}) must be a " - f"multiple of policy.automodel_cfg.context_parallel_size ({student_cp})." - ) + if not _megatron_enabled(policy_config): + student_cp = policy_config["automodel_cfg"]["context_parallel_size"] + assert student_div % student_cp == 0, ( + f"policy.make_sequence_length_divisible_by ({student_div}) must be a " + f"multiple of policy.automodel_cfg.context_parallel_size ({student_cp})." + ) for i, (teacher, tc) in enumerate(zip(teachers, teacher_configs)): + if _megatron_enabled(tc): + continue teacher_cp = tc["automodel_cfg"]["context_parallel_size"] - if teacher.aligner.projection_matrix_path is None: + if not _teacher_is_cross_tokenizer(master_config, teacher): # A same-vocab teacher forwards the student batch. assert student_div % teacher_cp == 0, ( f"teachers.{i} is a same-vocab teacher, which forwards the " @@ -347,7 +1682,7 @@ def setup( # no alignment — the loss does a direct per-position KL there). aligners: list[Optional[TokenAligner]] = [ None - if teacher.aligner.projection_matrix_path is None + if not _teacher_is_cross_tokenizer(master_config, teacher) else TokenAligner( student_tokenizer=student_tokenizer, teacher_tokenizer=teacher_tokenizers[i], @@ -356,7 +1691,26 @@ def setup( for i, teacher in enumerate(teachers) ] - collator = CrossTokenizerCollator( + collator_cls = ( + CSDCrossTokenizerCollator + if master_config.algorithm == "csd" + else CrossTokenizerCollator + ) + csd_collator_options: dict[str, bool] = {} + if master_config.algorithm == "csd": + # CSD never truncates one side alone: a text row that does not fit + # every side is clipped to a prefix that does, and overlength chat + # rows are rejected. Text rows carry each tokenizer's special + # tokens; ordinary chat supervises the final assistant turn's + # content + end-of-turn token with both CE and KD. + csd_collator_options = { + "require_routed_experts": router_replay_enabled(policy_config), + "reject_overlength_text": True, + } + if master_config.collator.mode == "text": + csd_collator_options["add_special_tokens_in_text"] = True + csd_collator_options["clip_overlength_text"] = True + collator = collator_cls( config=master_config.collator, student_tokenizer=student_tokenizer, teacher_tokenizers=list(teacher_tokenizers), @@ -370,6 +1724,7 @@ def setup( drop_first_assistant_chunk_kl_by_teacher=[ teacher.aligner.drop_first_assistant_chunk_kl for teacher in teachers ], + **csd_collator_options, ) # ========================== @@ -395,6 +1750,19 @@ def setup( flush=True, ) + # Megatron uses train_iters to configure its scheduler horizon. Derive it + # from the same bounds that terminate the training loop so it stays in sync + # with the actual dataloader rather than becoming a stale YAML constant. + total_train_iters = min( + distillation_config["max_num_steps"], + distillation_config["max_num_epochs"] * len(train_dataloader), + ) + for tc in teacher_configs: + if "megatron_cfg" in tc and tc["megatron_cfg"]["enabled"]: + tc["megatron_cfg"]["train_iters"] = total_train_iters + if "megatron_cfg" in policy_config and policy_config["megatron_cfg"]["enabled"]: + policy_config["megatron_cfg"]["train_iters"] = total_train_iters + val_dataloader: Optional[StatefulDataLoader] = None if val_dataset is not None and ( distillation_config["val_period"] > 0 @@ -420,7 +1788,7 @@ def setup( # ========================== print("\n▶ Setting up compute cluster...", flush=True) cluster = RayVirtualCluster( - name="xtoken_off_policy_distillation_cluster", + name=f"{master_config.algorithm}_off_policy_distillation_cluster", bundle_ct_per_node_list=[cluster_config.gpus_per_node] * cluster_config.num_nodes, use_gpus=True, @@ -463,6 +1831,16 @@ def setup( init_optimizer=True, init_reference_model=False, ) + # Finish the student's lazy worker/model construction while every teacher + # is still offloaded. Without this barrier, colocated Megatron student + # construction can overlap the first teacher forward and OOM at its peak. + # Megatron-FSDP cannot CPU-offload its shards in the current MCore version, + # but its steady-state resident shard is much smaller than construction. + student_policy.prepare_for_training() + if distillation_config.get("offload_student_after_step", False): + # Backends that support student CPU offload should start the first + # teacher phase in the same offloaded state used between later steps. + student_policy.offload_after_refit() # ========================== # Teacher/student grid @@ -472,8 +1850,14 @@ def setup( # order) and tile it cleanly into per-DP-rank chunks and whole microbatches. # assert_teacher_student_batch_grid checks both (GBS agreement + tiling). student_dp = student_policy.data_parallel_size - student_tp = policy_config["automodel_cfg"]["tensor_parallel_size"] - student_cp = policy_config["automodel_cfg"]["context_parallel_size"] + if configured_packing_dp is not None and student_dp != configured_packing_dp: + raise RuntimeError( + "CSD lockstep packing configured DP changed during worker setup: " + f"expected {configured_packing_dp}, student reports {student_dp}." + ) + student_tp, student_cp, student_pp = _xtoken_entity_parallelism( + policy_config, label="policy" + ) # Each teacher may differ from the student (and from each other) in # DP/MBS/TP/CP, so check the batch grid and node-local IPC layout per # teacher. Train and validation share the grid (the student reuses its train @@ -481,6 +1865,12 @@ def setup( # so one check per teacher covers both. for i, (teacher_policy, tc) in enumerate(zip(teacher_policies, teacher_configs)): teacher_dp = teacher_policy.data_parallel_size + if configured_packing_dp is not None and teacher_dp != configured_packing_dp: + raise RuntimeError( + "CSD lockstep packing configured DP changed during worker setup: " + f"expected {configured_packing_dp}, teachers[{i}] reports " + f"{teacher_dp}." + ) assert_teacher_student_batch_grid( global_batch_size=distillation_config["num_prompts_per_step"], student_gbs=policy_config["train_global_batch_size"], @@ -492,14 +1882,22 @@ def setup( ) # Node-local CUDA IPC: on >1 node it only works when teacher/student # share DP and a node-aligned model-parallel group, else a student rank - # would read teacher shards from another node. + # would read teacher shards from another node. The teacher's TP/CP/PP + # are read per backend (DTensor reports PP=1), and the model-parallel + # group is tp * cp * pp — PP counts because the producing last stage + # must land on the same node as the consuming student ranks. + teacher_tp, teacher_cp, teacher_pp = _xtoken_entity_parallelism( + tc, label=f"teachers[{i}]" + ) assert_xtoken_ipc_node_local( num_nodes=cluster_config.num_nodes, gpus_per_node=cluster_config.gpus_per_node, student_tp=student_tp, student_cp=student_cp, - teacher_tp=tc["automodel_cfg"]["tensor_parallel_size"], - teacher_cp=tc["automodel_cfg"]["context_parallel_size"], + student_pp=student_pp, + teacher_tp=teacher_tp, + teacher_cp=teacher_cp, + teacher_pp=teacher_pp, student_dp=student_dp, teacher_dp=teacher_dp, ) @@ -515,8 +1913,8 @@ def setup( # Per-teacher metadata is injected as parallel lists (one entry per # `teachers[i]`); the loss fn reads these lists directly. `len(tokenizer)` # is the HF embedding/lm_head dim, sizing each projection matrix's V_s/V_t - # axes exactly. No `send_full_logits`: every teacher ships full logits and - # the loss derives the top-k subset student-side. + # axes exactly. Dense IPC ships full logits; sparse IPC ships teacher top-k + # logits plus logZ and lets the loss reconstruct only the required support. loss_config = { **loss_config, "student_vocab_size": len(student_tokenizer), @@ -525,10 +1923,22 @@ def setup( teacher.aligner.projection_matrix_path for teacher in teachers ], "teacher_weights": [t.weight for t in teachers], - "teacher_gold_loss": [t.gold_loss for t in teachers], - "teacher_xtoken_loss": [t.xtoken_loss for t in teachers], + # CSD pseudo-target tables (student<->teacher sub-token chains) per + # cross-tokenizer teacher; None for same-tokenizer teachers. + "student_to_teacher_table_paths": [ + t.aligner.student_to_teacher_table_path for t in teachers + ], + "teacher_to_student_table_paths": [ + t.aligner.teacher_to_student_table_path for t in teachers + ], } - loss_fn = CrossTokenizerDistillationLossFn(loss_config) + if master_config.algorithm == "csd": + loss_config["owner_local"] = owner_local + loss_fn = CSDLossFn(loss_config) + else: + loss_config["teacher_gold_loss"] = [t.gold_loss for t in teachers] + loss_config["teacher_xtoken_loss"] = [t.xtoken_loss for t in teachers] + loss_fn = CrossTokenizerDistillationLossFn(loss_config) print("\n" + "=" * 60) print(" " * 18 + "SETUP COMPLETE") @@ -547,6 +1957,246 @@ def setup( ) +def _batch_length_tuple(batch: BatchedDataDict[Any], key: str) -> tuple[int, ...]: + value = batch.get(key) + if not torch.is_tensor(value) or value.ndim != 1: + raise ValueError( + f"CSD lockstep packing requires {key!r} as a rank-1 tensor of " + f"exact untruncated lengths, got {type(value).__name__}." + ) + lengths = tuple(int(length) for length in value.detach().cpu().tolist()) + if len(lengths) != batch.size or any(length <= 0 for length in lengths): + raise ValueError( + f"CSD lockstep packing received invalid {key}: {lengths}; expected " + f"{batch.size} positive logical-sample lengths." + ) + return lengths + + +def _effective_lengths( + raw_lengths: tuple[int, ...], *, divisor: int +) -> tuple[int, ...]: + divisor = _positive_config_int(divisor, field="sequence length divisor") + return tuple( + ((length + divisor - 1) // divisor) * divisor for length in raw_lengths + ) + + +def _fixed_tail_physical_size( + effective_lengths: tuple[int, ...], *, capacity: int +) -> int: + """Keep DTensor TP's fixed packed tail without hiding an oversized bin.""" + return max(capacity, sum(effective_lengths)) + + +def _bind_semantic_regions_to_occurrences( + batch: BatchedDataDict[Any], batch_item_ids: tuple[int, ...] +) -> None: + """Attach occurrence identity to side-local chat semantic-region records.""" + for key in tuple(batch.keys()): + if key != "student_semantic_regions" and not ( + key.startswith("teacher_") and key.endswith("_semantic_regions") + ): + continue + per_sample = batch[key] + if not isinstance(per_sample, Sequence) or isinstance(per_sample, (str, bytes)): + raise ValueError(f"{key} must contain one semantic-region list per row") + if len(per_sample) != len(batch_item_ids): + raise ValueError( + f"{key} has {len(per_sample)} rows for " + f"{len(batch_item_ids)} batch occurrences" + ) + bound_rows: list[tuple[tuple[object, ...], ...]] = [] + for batch_item_id, regions in zip(batch_item_ids, per_sample, strict=True): + bound_regions: list[tuple[object, ...]] = [] + for region in regions: + if not isinstance(region, Sequence) or isinstance(region, (str, bytes)): + raise ValueError(f"{key} contains a non-sequence region") + region_tuple = tuple(region) + if len(region_tuple) == 5: + bound_regions.append((batch_item_id, *region_tuple)) + elif len(region_tuple) == 6 and region_tuple[0] == batch_item_id: + bound_regions.append(region_tuple) + else: + raise ValueError( + f"{key} region must be (turn, role, name, start, end) " + "before binding or begin with its exact batch_item_id" + ) + bound_rows.append(tuple(bound_regions)) + batch[key] = bound_rows + + +def build_csd_lockstep_packing_plan( + batch: BatchedDataDict[Any], + master_config: MasterConfig, + *, + batch_uid: int, + data_parallel_size: int, +) -> Optional[LockstepPackingPlan]: + """Assign occurrence IDs and build the one controller-owned packing plan.""" + if not _sequence_packing_enabled(master_config.policy): + return None + if batch_uid < 0 or batch_uid >= 2**31: + raise ValueError( + f"batch_uid must fit the non-negative high 31 bits of int64, got {batch_uid}." + ) + + sample_ids_value = batch.get("sample_id") + if torch.is_tensor(sample_ids_value): + sample_ids: Sequence[object] = sample_ids_value.detach().cpu().tolist() + elif isinstance(sample_ids_value, Sequence) and not isinstance( + sample_ids_value, (str, bytes) + ): + sample_ids = sample_ids_value + else: + raise ValueError( + "CSD lockstep packing requires a durable sample_id for every " + "logical row before occurrence IDs are assigned." + ) + if len(sample_ids) != batch.size: + raise ValueError( + f"sample_id has {len(sample_ids)} entries for batch size {batch.size}." + ) + + batch_item_ids = tuple((batch_uid << 32) | ordinal for ordinal in range(batch.size)) + length_device = batch["input_lengths"].device + batch["batch_item_id"] = torch.tensor( + batch_item_ids, dtype=torch.long, device=length_device + ) + _bind_semantic_regions_to_occurrences(batch, batch_item_ids) + items = tuple( + LockstepPackingItem(sample_id=sample_id, batch_item_id=batch_item_id) + for sample_id, batch_item_id in zip(sample_ids, batch_item_ids, strict=True) + ) + + student_config = master_config.policy + student_raw_lengths = _batch_length_tuple(batch, "input_lengths") + student_divisor = int(student_config["make_sequence_length_divisible_by"]) + student_packing = student_config["sequence_packing"] + student_capacity = int(student_packing["train_mb_tokens"]) + student_physical_size_fn = None + if ( + student_config["automodel_cfg"]["enabled"] + and int(student_config["automodel_cfg"]["tensor_parallel_size"]) > 1 + ): + # Generic Automodel packing pads TP>1 training microbatches to the + # configured token budget. Preserve that requirement in the shared + # plan so controller materialization cannot bypass it. + student_physical_size_fn = partial( + _fixed_tail_physical_size, capacity=student_capacity + ) + side_specs = [ + SidePackingSpec( + side_id="student", + capacity=student_capacity, + raw_lengths=student_raw_lengths, + effective_lengths=_effective_lengths( + student_raw_lengths, divisor=student_divisor + ), + physical_size_fn=student_physical_size_fn, + ) + ] + for teacher_idx, teacher in enumerate(master_config.teachers): + teacher_config = teacher.policy_config() + if not _teacher_is_cross_tokenizer(master_config, teacher): + teacher_raw_lengths = student_raw_lengths + else: + teacher_raw_lengths = _batch_length_tuple( + batch, f"teacher_{teacher_idx}_input_lengths" + ) + teacher_packing = teacher_config["sequence_packing"] + side_specs.append( + SidePackingSpec( + side_id=f"teacher_{teacher_idx}", + capacity=int(teacher_packing["logprob_mb_tokens"]), + raw_lengths=teacher_raw_lengths, + effective_lengths=_effective_lengths( + teacher_raw_lengths, + divisor=int(teacher_config["make_sequence_length_divisible_by"]), + ), + ) + ) + + return build_lockstep_packing_plan( + batch_uid=batch_uid, + items=items, + sides=side_specs, + data_parallel_size=data_parallel_size, + ) + + +def log_csd_packing_telemetry( + plan: LockstepPackingPlan, batch: BatchedDataDict[Any] +) -> None: + """Emit controller-side packing evidence without conflating bins and samples.""" + sample_mask = batch["sample_mask"] + if torch.is_tensor(sample_mask): + masked_samples = int((sample_mask == 0).sum().item()) + else: + masked_samples = sum(float(value) == 0.0 for value in sample_mask) + multi_sample_bins = sum(len(bin_items) > 1 for bin_items in plan.bins) + split_count = plan.equalization_splits + print( + "CSD_PACKING " + f"batch_uid={plan.batch_uid} logical_samples=" + f"{len(plan.canonical_batch_item_ids)} physical_bins={len(plan.bins)} " + f"multi_sample_bins={multi_sample_bins} masked_logical_samples=" + f"{masked_samples} equalization_splits={split_count} transport=dense", + flush=True, + ) + for side_id, side in plan.sides.items(): + raw_tokens = sum(side.raw_lengths) + effective_tokens = sum(side.effective_lengths) + physical_tokens = sum(side.physical_tokens_by_bin) + available_tokens = side.capacity * len(side.physical_tokens_by_bin) + utilization = physical_tokens / available_tokens + print( + "CSD_PACKING_SIDE " + f"batch_uid={plan.batch_uid} side={side_id} bins_per_dp_rank=" + f"{tuple(len(indices) for indices in side.rank_bin_indices)} " + f"raw_tokens={raw_tokens} effective_tokens={effective_tokens} " + f"physical_tokens={physical_tokens} capacity_tokens={available_tokens} " + f"utilization={utilization:.6f} padding_waste=" + f"{physical_tokens - raw_tokens} tail_padding_waste=" + f"{physical_tokens - effective_tokens}", + flush=True, + ) + + +def _dense_ipc_logical_bytes(handles: Sequence[Mapping[str, Any]]) -> int: + """Return the logical dense-logit bytes addressable by aggregated handles.""" + total = 0 + for sample_index, sample in enumerate(handles): + shards = sample.get("teacher_shards") + if not isinstance(shards, Sequence) or isinstance(shards, (str, bytes)): + raise ValueError( + "Dense IPC telemetry requires teacher_shards for sample " + f"{sample_index}." + ) + for shard in shards: + if not isinstance(shard, Mapping): + raise TypeError( + "Dense IPC telemetry requires mapping shard records, got " + f"{type(shard).__name__}." + ) + actual_shape = shard.get("actual_shape") + dtype = shard.get("dtype") + if ( + not isinstance(actual_shape, (list, tuple, torch.Size)) + or len(actual_shape) != 2 + or not isinstance(dtype, torch.dtype) + ): + raise ValueError( + "Dense IPC telemetry requires a two-dimensional actual_shape " + f"and torch dtype, got {actual_shape!r}/{dtype!r}." + ) + total += ( + math.prod(int(value) for value in actual_shape) + * torch.empty((), dtype=dtype).element_size() + ) + return total + + # =============================================================================== # Train loop # =============================================================================== @@ -559,28 +2209,39 @@ def export_teacher_logits_and_pack( teacher_mbs: list[int], *, timer: Optional[Timer] = None, + packing_plan: Optional[LockstepPackingPlan] = None, ) -> BatchedDataDict[Any]: """Serially run each teacher's forward and pack the student ``train_data``. Teachers run one at a time (collocated): each is onloaded for inference, - forwarded, then offloaded. Every teacher ships full-vocab logits over CUDA - IPC (``teacher_{i}_full_logits_ipc``) — the loss derives the microbatch-global - top-k subset student-side, so there is no top-K transport. A same-vocab - teacher (``projection_matrix_paths[i] is None``) reuses the student tokens - and rides no extra alignment; a cross-tokenizer teacher's own tokenization - and ``alignment_{i}_*`` payload ride along (teacher-indexed). The persistent - IPC buffers stay resident on the teacher GPUs until released by the caller - after ``student.train``. Shared by the train loop and ``validate`` so the - forward+pack sequence can't drift between them. + forwarded, then offloaded. Cross-tokenizer teachers use sparse top-k + logZ + IPC when ``teacher_topk_ipc_k > 0`` and otherwise use full-vocab IPC; + same-vocab teachers always retain full logits for their direct-KL path. A + cross-tokenizer teacher's own tokenization and ``alignment_{i}_*`` payload + ride along (teacher-indexed). The persistent IPC buffers stay resident on + the teacher GPUs until released by the caller after ``student.train``. + Shared by the train loop and ``validate`` so the forward+pack sequence can't + drift between them. """ train_data: dict[str, Any] = { "input_ids": batch["input_ids"], "input_lengths": batch["input_lengths"], "token_mask": batch["token_mask"], + "kd_token_mask": batch.get("kd_token_mask", batch["token_mask"]), "sample_mask": batch["sample_mask"], } + if "routed_experts" in batch: + train_data["routed_experts"] = batch["routed_experts"] + if packing_plan is not None: + train_data["sample_id"] = batch["sample_id"] + train_data["batch_item_id"] = batch["batch_item_id"] + if "student_semantic_regions" in batch: + train_data["student_semantic_regions"] = batch["student_semantic_regions"] for i, teacher_policy in enumerate(teacher_policies): - same_vocab = loss_fn.projection_matrix_paths[i] is None + same_vocab = not _loss_teacher_is_cross_tokenizer(loss_fn, i) + semantic_key = f"teacher_{i}_semantic_regions" + if packing_plan is not None and semantic_key in batch: + train_data[semantic_key] = batch[semantic_key] if same_vocab: # Same tokenizer: the teacher forward reuses the student tokens. teacher_data: BatchedDataDict[Any] = BatchedDataDict( @@ -608,11 +2269,48 @@ def export_teacher_logits_and_pack( ): train_data[f"alignment_{i}_{field}"] = batch[f"alignment_{i}_{field}"] + if packing_plan is not None: + teacher_data["sample_id"] = batch["sample_id"] + teacher_data["batch_item_id"] = batch["batch_item_id"] + teacher_policy.prepare_for_lp_inference() - handles = teacher_policy.get_full_logits_ipc( - teacher_data, micro_batch_size=teacher_mbs[i], timer=timer + force_include_token_ids = None + if not same_vocab and isinstance(loss_fn, CSDLossFn): + if loss_fn.cfg.owner_local and loss_fn.cfg.teacher_topk_ipc_k > 0: + force_include_token_ids = _build_teacher_force_include_token_id_slots( + batch, teacher_idx=i, kl_chunk_shift=loss_fn.cfg.kl_chunk_shift + ) + else: + force_include_token_ids = _build_teacher_force_include_token_ids( + batch, + teacher_idx=i, + loss_config=loss_fn.cfg.model_dump(), + ) + ipc_suffix, handles = _get_teacher_logits_ipc( + teacher_policy, + teacher_data, + loss_fn, + teacher_idx=i, + micro_batch_size=teacher_mbs[i], + force_include_token_ids=force_include_token_ids, + timer=timer, + packing_plan=packing_plan, + packing_side_id=f"teacher_{i}" if packing_plan is not None else None, ) - train_data[f"teacher_{i}_full_logits_ipc"] = handles + train_data[f"teacher_{i}_{ipc_suffix}"] = handles + if packing_plan is not None: + if ipc_suffix != "full_logits_ipc": + raise ValueError( + "CSD lockstep telemetry encountered unsupported packed " + f"transport {ipc_suffix!r}." + ) + ipc_bytes = _dense_ipc_logical_bytes(handles) + print( + "CSD_IPC " + f"batch_uid={packing_plan.batch_uid} teacher=teacher_{i} " + f"transport=dense ipc_bytes={ipc_bytes}", + flush=True, + ) # Free the teacher's PARAMS to CPU; the persistent IPC buffers live in # worker state and survive this call. teacher_policy.offload_after_refit() @@ -620,6 +2318,44 @@ def export_teacher_logits_and_pack( return BatchedDataDict(train_data) +def _packing_batch_uids_from_state( + state: OffPolicyDistillationSaveState, +) -> Iterator[int]: + """Yield occurrence namespaces from the checkpointed high-water mark. + + A single persisted counter is shared by training and every validation batch. + This keeps ``batch_item_id`` values unique when validation contains several + batches and when the loop resumes from a checkpoint. Older checkpoints do + not carry the field, so they retain the former train-derived starting point. + """ + if "next_packing_batch_uid" not in state: + state["next_packing_batch_uid"] = int(state["total_steps"]) * 2 + while True: + batch_uid = state["next_packing_batch_uid"] + if isinstance(batch_uid, bool) or not isinstance(batch_uid, int): + raise ValueError( + f"next_packing_batch_uid must be an integer, got {batch_uid!r}." + ) + if batch_uid < 0 or batch_uid >= 2**31: + raise ValueError( + "next_packing_batch_uid must fit the non-negative high 31 bits " + f"of int64, got {batch_uid}." + ) + state["next_packing_batch_uid"] = batch_uid + 1 + yield batch_uid + + +def _restore_student_between_teacher_state( + student_policy: Policy, + master_config: MasterConfig, +) -> None: + """Restore the configured student residency between teacher forwards.""" + if master_config.distillation.get("offload_student_after_step", False): + student_policy.offload_after_refit() + elif master_config.policy.get("offload_optimizer_for_logprob", False): + student_policy.offload_before_refit() + + def xtoken_off_policy_distillation_train( student_policy: Policy, teacher_policies: list[Policy], @@ -650,6 +2386,24 @@ def xtoken_off_policy_distillation_train( val_at_end = distill_cfg["val_at_end"] max_epochs = distill_cfg["max_num_epochs"] max_steps = distill_cfg["max_num_steps"] + offload_student_after_step = bool( + distill_cfg.get("offload_student_after_step", False) + ) + offload_student_optimizer_after_step = bool( + master_config.policy.get("offload_optimizer_for_logprob", False) + ) + if offload_student_after_step: + print( + "Student post-step offload enabled: the student model/optimizer are " + "offloaded after each optimizer step, before the next teacher pass.", + flush=True, + ) + elif offload_student_optimizer_after_step: + print( + "Student post-step optimizer offload enabled: optimizer state is " + "offloaded before the next teacher pass while model shards remain resident.", + flush=True, + ) # Per-teacher export MBS (each teacher's own train MBS) and the # non-student-seq keys the worker's check_sequence_dim must skip # (teacher-count-dependent, so built from the loss fn). @@ -657,6 +2411,14 @@ def xtoken_off_policy_distillation_train( t.policy_config()["train_micro_batch_size"] for t in master_config.teachers ] skip_keys = xtoken_non_student_seq_keys(loss_fn) + # Training and validation draw from one checkpointed high-water mark. A + # step-derived train counter is insufficient because one validation pass can + # consume many global batches before the next checkpoint/resume boundary. + packing_batch_uids = ( + _packing_batch_uids_from_state(off_policy_distillation_state) + if master_config.algorithm == "csd" + else None + ) if val_at_start and total_steps == 0 and val_dataloader is not None: val_metrics, val_timings = validate( @@ -667,6 +2429,7 @@ def xtoken_off_policy_distillation_train( master_config, skip_keys=skip_keys, timer=timer, + packing_batch_uids=packing_batch_uids, ) logger.log_metrics(val_metrics, total_steps, prefix="validation") logger.log_metrics(val_timings, total_steps, prefix="timing/validation") @@ -686,14 +2449,33 @@ def xtoken_off_policy_distillation_train( ) maybe_gpu_profile_step(student_policy, total_steps + 1) + packing_plan = None + if master_config.algorithm == "csd": + assert packing_batch_uids is not None + batch_uid = next(packing_batch_uids) + log_csd_logical_batch_digest(batch, batch_uid=batch_uid) + packing_plan = build_csd_lockstep_packing_plan( + batch, + master_config, + batch_uid=batch_uid, + data_parallel_size=student_policy.data_parallel_size, + ) + if packing_plan is not None: + log_csd_packing_telemetry(packing_plan, batch) + with timer.time("total_step_time"): with timer.time("teacher_forward"): - # Serial per-teacher forward; each teacher ships full-vocab - # logits over IPC (``teacher_{i}_full_logits_ipc``) and the - # loss derives the microbatch-global top-k subset - # student-side. Packs the per-teacher alignment payload too. + # Serial per-teacher forward; each cross-tokenizer teacher + # selects dense or sparse IPC from the loss config while a + # same-vocab teacher keeps full logits. The per-teacher + # alignment payload is packed alongside the handles. train_data = export_teacher_logits_and_pack( - teacher_policies, loss_fn, batch, teacher_mbs, timer=timer + teacher_policies, + loss_fn, + batch, + teacher_mbs, + timer=timer, + packing_plan=packing_plan, ) with timer.time("training_prep"): @@ -701,12 +2483,25 @@ def xtoken_off_policy_distillation_train( with timer.time("policy_training"): try: + packing_kwargs: dict[str, Any] = {} + if packing_plan is not None: + packing_kwargs = { + "packing_plan": packing_plan, + "packing_side_id": "student", + } train_results = student_policy.train( train_data, loss_fn, timer=timer, check_dim_skip_keys=skip_keys, + **packing_kwargs, ) + if packing_plan is not None: + _log_dense_ipc_reconstruction_telemetry( + train_results.get("all_mb_metrics", {}), + packing_plan=packing_plan, + num_teachers=loss_fn.num_teachers, + ) except Exception: # Free every teacher's producer IPC buffer before # propagating so a failed step doesn't leak teacher @@ -718,6 +2513,17 @@ def xtoken_off_policy_distillation_train( teacher_policy.release_ipc_buffer() raise + if offload_student_after_step: + with timer.time("student_step_offload"): + _restore_student_between_teacher_state( + student_policy, master_config + ) + elif offload_student_optimizer_after_step: + with timer.time("student_optimizer_offload"): + _restore_student_between_teacher_state( + student_policy, master_config + ) + is_last_step = (total_steps + 1 >= max_steps) or ( (current_epoch + 1 == max_epochs) and (current_step + 1 == len(dataloader)) @@ -736,6 +2542,7 @@ def xtoken_off_policy_distillation_train( master_config, skip_keys=skip_keys, timer=timer, + packing_batch_uids=packing_batch_uids, ) logger.log_metrics( val_metrics, total_steps + 1, prefix="validation" @@ -751,23 +2558,11 @@ def xtoken_off_policy_distillation_train( metrics.update(train_results["all_mb_metrics"]) # Reduce per-microbatch metrics to per-step scalars. The # P-KL path emits kl_loss/ce_loss/kl_loss_scale/proj_accuracy; - # the gold-loss path emits kl_common/l1_uncommon. Either set - # may be present — reduce both via the same rules. + # the gold-loss path emits kl_common/l1_uncommon; the CSD path + # emits the *_per_chunk diagnostics. Any set may be present — + # reduce all via the same rules. for k, v in metrics.items(): - if k in { - "lr", - "wd", - "global_valid_seqs", - "global_valid_toks", - "accuracy", - "proj_accuracy", - "kl_loss_scale", - "kl_common", - "l1_uncommon", - }: - metrics[k] = float(np.mean(v)) - else: - metrics[k] = float(np.sum(v)) + metrics[k] = reduce_mb_metric(k, v) if "global_valid_toks" in metrics: total_valid_tokens += int(metrics["global_valid_toks"]) @@ -838,6 +2633,9 @@ def xtoken_off_policy_distillation_train( ckpt_path, wait_fn=student_policy.finalize_async_save, ) + _restore_student_between_teacher_state( + student_policy, master_config + ) # ===== Logging ===== timing_metrics: dict[str, float] = timer.get_timing_metrics( @@ -854,6 +2652,7 @@ def xtoken_off_policy_distillation_train( f"(per-MB-mean: {metrics['loss'] / n_mb:.4f})", flush=True, ) + print(f" • GradNorm: {metrics['grad_norm']:.4f}", flush=True) # P-KL path metrics — only printed when they're present. if "kl_loss" in metrics: kl_sum = float(metrics["kl_loss"]) @@ -947,6 +2746,7 @@ def validate( master_config: MasterConfig, skip_keys: frozenset[str], timer: Optional[Timer] = None, + packing_batch_uids: Optional[Iterator[int]] = None, ) -> tuple[dict[str, Any], dict[str, Any]]: """Held-out KL/CE on a validation dataloader. @@ -959,6 +2759,8 @@ def validate( """ distill_cfg = master_config.distillation timer = timer if timer is not None else Timer() + if master_config.algorithm == "csd" and packing_batch_uids is None: + packing_batch_uids = count() losses: list[float] = [] # The P-KL path emits kl_loss/ce_loss; the gold path emits @@ -991,18 +2793,43 @@ def validate( for batch in val_dataloader: target_size = math.ceil(batch.size / pad_quantum) * pad_quantum batch = pad_distillation_val_batch(batch, target_size) + packing_plan = None + if master_config.algorithm == "csd": + assert packing_batch_uids is not None + packing_plan = build_csd_lockstep_packing_plan( + batch, + master_config, + batch_uid=next(packing_batch_uids), + data_parallel_size=student_policy.data_parallel_size, + ) + if packing_plan is not None: + log_csd_packing_telemetry(packing_plan, batch) train_data = export_teacher_logits_and_pack( - teacher_policies, loss_fn, batch, teacher_mbs, timer=timer + teacher_policies, + loss_fn, + batch, + teacher_mbs, + timer=timer, + packing_plan=packing_plan, ) student_policy.prepare_for_training() try: + packing_kwargs: dict[str, Any] = {} + if packing_plan is not None: + packing_kwargs = { + "packing_plan": packing_plan, + "packing_side_id": "student", + } results = student_policy.train( train_data, loss_fn, eval_mode=True, + gbs=target_size, check_dim_skip_keys=skip_keys, + **packing_kwargs, ) + _restore_student_between_teacher_state(student_policy, master_config) except Exception: for teacher_policy in teacher_policies: teacher_policy.release_ipc_buffer() diff --git a/nemo_rl/data/__init__.py b/nemo_rl/data/__init__.py index 0638c7abf7a..67a547148cf 100644 --- a/nemo_rl/data/__init__.py +++ b/nemo_rl/data/__init__.py @@ -51,6 +51,16 @@ class ResponseDatasetConfig(TypedDict): pad_audio_to_hop_length: NotRequired[bool] min_generation_tokens: NotRequired[int] max_samples: NotRequired[int | None] + # Arrow/raw-text dataset fields used by cross-tokenizer distillation. + data_files: NotRequired[str | list[str] | None] + text_key: NotRequired[str | None] + characters_per_sample: NotRequired[int | None] + # Native OpenAI-format chat adapter fields. + chat_key: NotRequired[str] + use_preserving_dataset: NotRequired[bool] + system_key: NotRequired[str | None] + system_prompt: NotRequired[str | None] + tool_key: NotRequired[str | None] class PreferenceDatasetConfig(TypedDict): diff --git a/nemo_rl/data/cross_tokenizer/csd_collate.py b/nemo_rl/data/cross_tokenizer/csd_collate.py new file mode 100644 index 00000000000..5b6dc3ae233 --- /dev/null +++ b/nemo_rl/data/cross_tokenizer/csd_collate.py @@ -0,0 +1,349 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CSD collation: shared cross-tokenizer batches plus durable row identities.""" + +from __future__ import annotations + +from collections.abc import MutableMapping +from typing import Any, List, Optional, cast + +import torch +from transformers.tokenization_utils_base import PreTrainedTokenizerBase + +from nemo_rl.algorithms.x_token.token_aligner import TokenAligner +from nemo_rl.data.cross_tokenizer.cross_tokenizer_collate import ( + CrossTokenizerCollator, + CrossTokenizerCollatorConfig, +) +from nemo_rl.data.interfaces import DatumSpec +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +class CSDCrossTokenizerCollator(CrossTokenizerCollator): + """Shared cross-tokenizer collation plus CSD row identities and MoE routes. + + Batches are built by :class:`CrossTokenizerCollator`. Every batch also + carries ``sample_id``, the durable source identity of each row that CSD's + lockstep packing keys on, and, with ``require_routed_experts``, + ``routed_experts``: the rollout-recorded MoE routes of every student + token, right-padded to ``[B, T, num_layers, topk]``. + + Args: + require_routed_experts: Emit ``routed_experts`` from each message's + rollout-recorded ``token_ids`` and ``routed_experts``, which must + reproduce the student tokenization exactly. + clip_overlength_text: Text-mode policy for rows that do not fit. A + row fits when, on the student and on every teacher, its + tokenization rounded up to that side's sequence-length divisor + is within that side's context length; same-tokenizer teachers + use the student tokenization. ``True`` replaces the text of each + row that does not fit with a prefix that fits every side, so the + student and every teacher still cover the same text. Requires + ``mode="text"``. + All other arguments are those of :class:`CrossTokenizerCollator`. + """ + + def __init__( + self, + *, + config: CrossTokenizerCollatorConfig, + student_tokenizer: PreTrainedTokenizerBase, + teacher_tokenizers: List[Optional[PreTrainedTokenizerBase]], + aligners: List[Optional[TokenAligner]], + ctx_length_student: int, + ctx_length_teachers: List[int], + drop_first_assistant_chunk_kl_by_teacher: List[bool], + require_routed_experts: bool, + make_seq_div_by_student: int = 1, + make_seq_div_by_teachers: Optional[List[int]] = None, + reject_overlength_text: bool = False, + ce_covers_assistant_turn: bool = False, + add_special_tokens_in_text: bool = False, + clip_overlength_text: bool = False, + ) -> None: + if clip_overlength_text and config.mode != "text": + raise ValueError( + "clip_overlength_text requires mode='text'; chat rows are not clipped." + ) + super().__init__( + config=config, + student_tokenizer=student_tokenizer, + teacher_tokenizers=teacher_tokenizers, + aligners=aligners, + ctx_length_student=ctx_length_student, + ctx_length_teachers=ctx_length_teachers, + drop_first_assistant_chunk_kl_by_teacher=drop_first_assistant_chunk_kl_by_teacher, + make_seq_div_by_student=make_seq_div_by_student, + make_seq_div_by_teachers=make_seq_div_by_teachers, + reject_overlength_text=reject_overlength_text, + ce_covers_assistant_turn=ce_covers_assistant_turn, + add_special_tokens_in_text=add_special_tokens_in_text, + ) + self.require_routed_experts = require_routed_experts + self.clip_overlength_text = clip_overlength_text + + def __call__(self, batch: List[DatumSpec]) -> BatchedDataDict[Any]: + sample_ids = self._required_sample_ids(batch) + if self.clip_overlength_text: + batch = self._clip_overlength_rows(batch) + out = super().__call__(batch) + out["sample_id"] = sample_ids + self._add_router_replay_metadata( + out, + batch, + student_input_ids=out["input_ids"], + student_input_lengths=out["input_lengths"], + ) + return out + + def _clip_overlength_rows(self, batch: List[DatumSpec]) -> List[DatumSpec]: + """``batch`` with each text row that does not fit every side clipped. + + Rows that fit are returned as is; clipped rows are copies whose text is + a prefix that fits every side. Non-string content is left for the + shared collator to reject. + """ + texts = [datum["message_log"][0]["content"] for datum in batch] + if not all(isinstance(text, str) for text in texts): + return batch + texts = cast(List[str], texts) + clipped: List[DatumSpec] = [] + for datum, text, fits in zip(batch, texts, self._fits_every_side(texts)): + if fits: + clipped.append(datum) + continue + message_log = cast(List[dict[str, Any]], datum["message_log"]) + prefix = self._fitting_text_prefix(text, datum["idx"]) + clipped.append( + cast( + DatumSpec, + { + **datum, + "message_log": [ + {**message_log[0], "content": prefix}, + *message_log[1:], + ], + }, + ) + ) + return clipped + + def _fits_every_side(self, texts: List[str]) -> List[bool]: + """Whether each text fits the student and every teacher.""" + + def lengths(tokenizer: PreTrainedTokenizerBase) -> List[int]: + encoded = tokenizer( + texts, + truncation=False, + add_special_tokens=self.add_special_tokens_in_text, + ) + return [len(ids) for ids in encoded["input_ids"]] + + student_lengths = lengths(self.student_tokenizer) + sides = [ + (student_lengths, self.ctx_length_student, self.make_seq_div_by_student) + ] + for i, aligner in enumerate(self.aligners): + teacher_tokenizer = self.teacher_tokenizers[i] + if aligner is None or teacher_tokenizer is None: + # Same-tokenizer teachers reuse the student tokenization. + teacher_lengths = student_lengths + else: + teacher_lengths = lengths(teacher_tokenizer) + sides.append( + ( + teacher_lengths, + self.ctx_length_teachers[i], + self.make_seq_div_by_teachers[i], + ) + ) + return [ + all( + self._padded_width(side_lengths[row], make_seq_div_by) <= ctx_length + for side_lengths, ctx_length, make_seq_div_by in sides + ) + for row in range(len(texts)) + ] + + def _fitting_text_prefix(self, text: str, row_idx: int) -> str: + """A non-empty prefix of ``text`` that fits every side. + + Binary search over the prefix length. Token counts need not grow + monotonically with the prefix length, so the search keeps only + prefixes whose encodings fit: the result always fits but need not be + the longest prefix that does. + """ + low, high, best = 1, len(text) - 1, "" + while low <= high: + middle = (low + high) // 2 + if self._fits_every_side([text[:middle]])[0]: + best = text[:middle] + low = middle + 1 + else: + high = middle - 1 + if not best: + raise ValueError( + f"sample idx={row_idx}: no non-empty prefix of the text fits " + "every context length" + ) + return best + + def _add_router_replay_metadata( + self, + out: MutableMapping[str, Any], + batch: List[DatumSpec], + *, + student_input_ids: torch.Tensor, + student_input_lengths: torch.Tensor, + ) -> None: + """Validate and batch rollout-recorded routes for the student forward. + + Router replay is meaningful only when the recorded routes describe the + exact token sequence that the student will train on. CSD normally + re-tokenizes source text/chat in the collator, so compare the persisted + per-message ``token_ids`` with that result before carrying the matching + ``routed_experts`` tensor forward. This turns stale tokenizer/template + settings into an actionable input error instead of silently replaying + routes at the wrong token positions. + """ + if not self.require_routed_experts: + return + + padded_seq_len = student_input_ids.shape[1] + routed_rows: list[torch.Tensor] = [] + route_shape: Optional[tuple[int, int]] = None + route_dtype: Optional[torch.dtype] = None + + for row_index, datum in enumerate(batch): + sample_id = datum["sample_id"] + token_parts: list[torch.Tensor] = [] + route_parts: list[torch.Tensor] = [] + message_log = cast(List[dict[str, Any]], datum["message_log"]) + for turn_index, message in enumerate(message_log): + token_ids_value = message.get("token_ids") + routed_experts_value = message.get("routed_experts") + if token_ids_value is None: + raise RuntimeError( + "policy.router_replay.enabled=true requires rollout-recorded " + "token_ids and routed_experts on every message; " + f"sample_id={sample_id!r}, turn_index={turn_index} is " + "missing metadata. Use a vLLM router-replay rollout as " + "the CSD data source." + ) + + token_ids = torch.as_tensor(token_ids_value, dtype=torch.long) + if token_ids.dim() != 1: + raise ValueError( + "router-replay token_ids must have shape [tokens]; " + f"sample_id={sample_id!r}, turn_index={turn_index}, " + f"got {tuple(token_ids.shape)}." + ) + # Some message serializers erase the trailing dimensions of an + # empty [0, L, K] tensor. Empty turns contribute no positions, + # so they can be skipped without weakening route coverage. + if token_ids.numel() == 0: + continue + if routed_experts_value is None: + raise RuntimeError( + "policy.router_replay.enabled=true requires rollout-recorded " + "token_ids and routed_experts on every non-empty message; " + f"sample_id={sample_id!r}, turn_index={turn_index} is " + "missing routed_experts. Use a vLLM router-replay " + "rollout as the CSD data source." + ) + routed_experts = torch.as_tensor(routed_experts_value) + if routed_experts.dim() != 3: + raise ValueError( + "router-replay routed_experts must have shape " + "[tokens, layers, topk]; " + f"sample_id={sample_id!r}, turn_index={turn_index}, " + f"got {tuple(routed_experts.shape)}." + ) + if routed_experts.shape[0] != token_ids.shape[0]: + raise ValueError( + "router-replay token_ids and routed_experts token axes " + f"differ for sample_id={sample_id!r}, " + f"turn_index={turn_index}: {token_ids.shape[0]} != " + f"{routed_experts.shape[0]}." + ) + + current_shape = ( + int(routed_experts.shape[1]), + int(routed_experts.shape[2]), + ) + if route_shape is None: + route_shape = current_shape + route_dtype = routed_experts.dtype + elif ( + current_shape != route_shape or routed_experts.dtype != route_dtype + ): + raise ValueError( + "router-replay routed_experts must use one [layers, topk] " + "shape and dtype across the batch; " + f"expected {route_shape}/{route_dtype}, got " + f"{current_shape}/{routed_experts.dtype} for " + f"sample_id={sample_id!r}, turn_index={turn_index}." + ) + token_parts.append(token_ids.detach().cpu()) + route_parts.append(routed_experts.detach().cpu()) + + if not token_parts or route_shape is None or route_dtype is None: + raise RuntimeError( + "policy.router_replay.enabled=true requires non-empty " + f"rollout route metadata; sample_id={sample_id!r} has none." + ) + + recorded_token_ids = torch.cat(token_parts, dim=0) + recorded_routes = torch.cat(route_parts, dim=0) + student_length = int(student_input_lengths[row_index].item()) + expected_token_ids = student_input_ids[row_index, :student_length].cpu() + if not torch.equal(recorded_token_ids, expected_token_ids): + raise RuntimeError( + "Cannot replay routed experts because rollout-recorded " + "student token_ids do not match the CSD student " + f"tokenization for sample_id={sample_id!r} " + f"(recorded={recorded_token_ids.shape[0]} tokens, " + f"retokenized={student_length}). Keep the rollout and " + "training tokenizer/chat-template settings identical." + ) + + pad_len = padded_seq_len - recorded_routes.shape[0] + if pad_len < 0: + raise RuntimeError( + "router-replay route metadata is longer than the student " + f"batch row for sample_id={sample_id!r}: " + f"{recorded_routes.shape[0]} > {padded_seq_len}." + ) + if pad_len: + recorded_routes = torch.nn.functional.pad( + recorded_routes, (0, 0, 0, 0, 0, pad_len), value=0 + ) + routed_rows.append(recorded_routes) + + out["routed_experts"] = torch.stack(routed_rows, dim=0) + + @staticmethod + def _required_sample_ids(batch: List[DatumSpec]) -> list[str | int]: + """Return durable source IDs; post-transform indices are not identity.""" + sample_ids: list[str | int] = [] + for row_index, datum in enumerate(batch): + sample_id = datum.get("sample_id") + if sample_id is None or not str(sample_id).strip(): + raise ValueError( + "CSDCrossTokenizerCollator requires a durable, non-empty " + "sample_id on every row; positional idx cannot survive " + f"filtering/splitting (batch row {row_index})." + ) + sample_ids.append(sample_id) + return sample_ids diff --git a/nemo_rl/data/datasets/response_datasets/arrow_text_dataset.py b/nemo_rl/data/datasets/response_datasets/arrow_text_dataset.py index c6b5d3b62a4..af2aeae5d8e 100644 --- a/nemo_rl/data/datasets/response_datasets/arrow_text_dataset.py +++ b/nemo_rl/data/datasets/response_datasets/arrow_text_dataset.py @@ -27,6 +27,9 @@ from datasets import Dataset from nemo_rl.data.datasets.raw_dataset import RawDataset +from nemo_rl.data.datasets.response_datasets.response_dataset import ( + attach_source_sample_id, +) from nemo_rl.data.datasets.utils import load_dataset_from_path @@ -70,8 +73,22 @@ def __init__( ): self.text_key = text_key self.task_name = "x_token" + self.source_namespace = (data_files, subset, split) raw = load_dataset_from_path(data_files, subset, split) + # Identity must be attached before filtering, mapping, or splitting so + # the surviving row keeps its source provenance. Respect an upstream + # ID when present; otherwise namespace the stable raw-row ordinal. + raw = raw.map( + attach_source_sample_id, + with_indices=True, + remove_columns=["sample_id"] if "sample_id" in raw.column_names else None, + fn_kwargs={ + "data_path": data_files, + "subset": subset, + "split": split, + }, + ) # Filter at the source so the packed and non-packed branches see the # same corpus (the packed path also drops empty/non-string rows). raw = raw.filter(lambda d: isinstance(d[text_key], str) and bool(d[text_key])) @@ -91,7 +108,7 @@ def __init__( "task_name": self.task_name, # Part of the HF datasets cache fingerprint; bump it # whenever the emitted row schema changes. - "schema_version": "messages-v1", + "schema_version": "messages-v2-sample-id", }, ) @@ -105,6 +122,7 @@ def format_data(self, data: dict[str, Any]) -> dict[str, Any]: return { "messages": [{"role": "assistant", "content": text}], "task_name": self.task_name, + "sample_id": data["sample_id"], } @@ -123,22 +141,27 @@ def _pack_generator( """ del schema_version buf: list[str] = [] + sample_ids: list[str] = [] n = 0 - for row in raw: + for source_ordinal, row in enumerate(raw): text = row[text_key] buf.append(text) + sample_ids.append(str(row.get("sample_id", source_ordinal))) n += len(text) if n >= characters_per_sample: packed = "\n".join(buf) yield { "messages": [{"role": "assistant", "content": packed}], "task_name": task_name, + "sample_id": "legacy-pack:" + "|".join(sample_ids), } buf = [] + sample_ids = [] n = 0 if buf: packed = "\n".join(buf) yield { "messages": [{"role": "assistant", "content": packed}], "task_name": task_name, + "sample_id": "legacy-pack:" + "|".join(sample_ids), } diff --git a/nemo_rl/data/datasets/response_datasets/oai_format_dataset.py b/nemo_rl/data/datasets/response_datasets/oai_format_dataset.py index 344c8aedeaf..b90ea90ddd8 100644 --- a/nemo_rl/data/datasets/response_datasets/oai_format_dataset.py +++ b/nemo_rl/data/datasets/response_datasets/oai_format_dataset.py @@ -19,6 +19,9 @@ from datasets import load_dataset from nemo_rl.data.datasets.raw_dataset import RawDataset +from nemo_rl.data.datasets.response_datasets.response_dataset import ( + attach_source_sample_id, +) class PreservingDataset: @@ -107,6 +110,8 @@ class OpenAIFormatDataset(RawDataset): heterogeneous schemas (e.g., for tool calls with varying argument structures). If False, uses standard HuggingFace dataset loading. Default is False for backward compatibility. + subset: Optional source subset/config used in synthesized sample IDs. + split: Source split used in synthesized sample IDs. Notes: - system_key and system_prompt are optional. If provided, it will be added @@ -127,6 +132,8 @@ def __init__( system_prompt: str | None = None, tool_key: str | None = "tools", use_preserving_dataset: bool = False, + subset: str | None = None, + split: str | None = "train", **kwargs, ): self.chat_key = chat_key @@ -141,6 +148,18 @@ def __init__( if not use_preserving_dataset: # Use the standard HuggingFace approach (faster and more standard) original_dataset = load_dataset("json", data_files=data_path)["train"] + original_dataset = original_dataset.map( + attach_source_sample_id, + with_indices=True, + remove_columns=["sample_id"] + if "sample_id" in original_dataset.column_names + else None, + fn_kwargs={ + "data_path": data_path, + "subset": subset, + "split": split, + }, + ) # Format the dataset self.dataset = original_dataset.map(self.format_data) @@ -175,8 +194,21 @@ def __init__( # Load JSON files directly with open(data_path, "r") as f: original_dataset = [json.loads(line) for line in f] + identified_dataset = [ + { + **item, + **attach_source_sample_id( + item, + raw_ordinal, + data_path=data_path, + subset=subset, + split=split, + ), + } + for raw_ordinal, item in enumerate(original_dataset) + ] # Format the dataset - formatted_data = [self.format_data(item) for item in original_dataset] + formatted_data = [self.format_data(item) for item in identified_dataset] # Use PreservingDataset to maintain exact structure self.dataset = PreservingDataset(formatted_data) @@ -201,6 +233,8 @@ def format_data(self, data: dict[str, Any]) -> dict[str, Any]: result = {"messages": messages, "task_name": self.task_name} if message_loss_mask is not None: result["message_loss_mask"] = message_loss_mask + if "sample_id" in data: + result["sample_id"] = data["sample_id"] if self.tool_key and self.tool_key in data: result["tools"] = data[self.tool_key] diff --git a/nemo_rl/data/datasets/response_datasets/response_dataset.py b/nemo_rl/data/datasets/response_datasets/response_dataset.py index e6197a501f2..427de2f7025 100644 --- a/nemo_rl/data/datasets/response_datasets/response_dataset.py +++ b/nemo_rl/data/datasets/response_datasets/response_dataset.py @@ -12,11 +12,83 @@ # See the License for the specific language governing permissions and # limitations under the License. +import hashlib +import json +from collections.abc import Mapping from typing import Any, Optional from nemo_rl.data.datasets.raw_dataset import RawDataset from nemo_rl.data.datasets.utils import load_dataset_from_path +_SOURCE_ID_KEYS = ("sample_id", "id", "uuid") + + +def resolve_source_sample_id( + data: Mapping[str, Any], + raw_ordinal: int, + *, + data_path: str, + subset: str | None, + split: str | None, +) -> str: + """Resolve a durable identity before dataset transformations. + + Explicit source identities are retained verbatim (after conversion to a + string). Otherwise, the identity combines a collision-resistant digest of + the complete dataset namespace with the original row ordinal. JSON + serialization keeps path/subset/split boundaries unambiguous. + + Args: + data: Unmodified source row. + raw_ordinal: Row position in the source before filtering or splitting. + data_path: Local path, shard/glob, URL, or Hugging Face dataset name. + subset: Optional Hugging Face subset/config name. + split: Requested source split. ``None`` has the loader's effective + default of ``"train"``. + + Returns: + The upstream identity, when present, or a deterministic synthesized ID. + """ + for key in _SOURCE_ID_KEYS: + value = data.get(key) + if value is not None: + source_id = str(value) + if source_id.strip(): + return source_id + + namespace = json.dumps( + { + "data_path": data_path, + "split": split or "train", + "subset": subset, + }, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ) + namespace_digest = hashlib.sha256(namespace.encode("utf-8")).hexdigest() + return f"source-row:{namespace_digest}:{raw_ordinal}" + + +def attach_source_sample_id( + data: Mapping[str, Any], + raw_ordinal: int, + *, + data_path: str, + subset: str | None, + split: str | None, +) -> dict[str, str]: + """Build the ``Dataset.map`` update that attaches source identity.""" + return { + "sample_id": resolve_source_sample_id( + data, + raw_ordinal, + data_path=data_path, + subset=subset, + split=split, + ) + } + class ResponseDataset(RawDataset): """Dataset class for response data which can be loaded from a JSON file. @@ -60,6 +132,22 @@ def __init__( # load from local or huggingface self.dataset = load_dataset_from_path(data_path, subset, split) + # Attach identity to untouched source rows. Formatting removes source + # columns and validation splitting reorders/selects rows, so neither + # operation may be allowed to invent identity afterward. + self.dataset = self.dataset.map( + attach_source_sample_id, + with_indices=True, + remove_columns=["sample_id"] + if "sample_id" in self.dataset.column_names + else None, + fn_kwargs={ + "data_path": data_path, + "subset": subset, + "split": split, + }, + ) + # format the dataset if "messages" not in self.dataset.column_names: self.dataset = self.dataset.map( @@ -82,4 +170,5 @@ def format_data(self, data: dict[str, Any]) -> dict[str, Any]: {"role": "assistant", "content": data[self.output_key]}, ], "task_name": self.task_name, + "sample_id": data["sample_id"], } diff --git a/nemo_rl/data/interfaces.py b/nemo_rl/data/interfaces.py index cce16869778..ff1bb033940 100644 --- a/nemo_rl/data/interfaces.py +++ b/nemo_rl/data/interfaces.py @@ -66,6 +66,14 @@ class DatumSpec(TypedDict): extra_env_info: Optional[dict[str, Any]] loss_multiplier: float # multiplier for the loss for this datum. 0 to mask out (say the sample is invalid) idx: int + # Durable identity of the source row. Unlike ``idx`` this is assigned + # before filtering/splitting and survives validation padding. + sample_id: NotRequired[str | int] + # Unique identity of one occurrence in a controller-owned global batch. + # CSD assigns this after validation padding and before lockstep packing. + batch_item_id: NotRequired[int] + # Canonical row-level tool schema used by tool-aware chat templates. + tools: NotRequired[list[dict[str, Any]]] task_name: NotRequired[str] # One binary value per message; assistant turns only. message_loss_mask: NotRequired[list[int]] diff --git a/nemo_rl/data/packing/__init__.py b/nemo_rl/data/packing/__init__.py index 9e539be501b..566268d88dc 100644 --- a/nemo_rl/data/packing/__init__.py +++ b/nemo_rl/data/packing/__init__.py @@ -23,6 +23,13 @@ SequencePacker, get_packer, ) +from nemo_rl.data.packing.lockstep import ( + LockstepPackingItem, + LockstepPackingPlan, + SidePackingPlan, + SidePackingSpec, + build_lockstep_packing_plan, +) from nemo_rl.data.packing.metrics import PackingMetrics __all__ = [ @@ -36,4 +43,9 @@ "ModifiedFirstFitDecreasingPacker", "get_packer", "PackingMetrics", + "LockstepPackingItem", + "SidePackingSpec", + "SidePackingPlan", + "LockstepPackingPlan", + "build_lockstep_packing_plan", ] diff --git a/nemo_rl/data/packing/lockstep.py b/nemo_rl/data/packing/lockstep.py new file mode 100644 index 00000000000..d7cf06ead91 --- /dev/null +++ b/nemo_rl/data/packing/lockstep.py @@ -0,0 +1,675 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Deterministic controller-side packing shared by multiple model sides. + +The planner in this module owns logical bin membership. Model-specific code may +materialize the geometry stored in :class:`SidePackingPlan`, but it must not +change ``LockstepPackingPlan.bins`` or independently repack a side. +""" + +from __future__ import annotations + +from collections.abc import Callable, Hashable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from fractions import Fraction +from numbers import Integral + +PhysicalSizeFn = Callable[[tuple[int, ...]], int] + + +def _normalize_int(value: object, *, name: str, minimum: int = 0) -> int: + if isinstance(value, bool) or not isinstance(value, Integral): + raise ValueError(f"{name} must be an integer, got {value!r}.") + normalized = int(value) + if normalized < minimum: + raise ValueError(f"{name} must be >= {minimum}, got {normalized}.") + return normalized + + +def _normalize_int_tuple( + values: Sequence[int], *, name: str, minimum: int = 0 +) -> tuple[int, ...]: + return tuple( + _normalize_int(value, name=f"{name}[{i}]", minimum=minimum) + for i, value in enumerate(values) + ) + + +def _cumulative_lengths(lengths: Sequence[int]) -> tuple[int, ...]: + result = [0] + for length in lengths: + result.append(result[-1] + length) + return tuple(result) + + +@dataclass(frozen=True) +class LockstepPackingItem: + """One occurrence in the canonical logical global batch. + + ``sample_id`` identifies the durable source row and may legitimately repeat, + for example after validation padding. ``batch_item_id`` identifies this one + occurrence and must be unique within the plan. + """ + + sample_id: Hashable + batch_item_id: int + + def __post_init__(self) -> None: + if self.sample_id is None: + raise ValueError("sample_id must not be None.") + try: + hash(self.sample_id) + except TypeError as error: + raise ValueError( + f"sample_id must be hashable, got {self.sample_id!r}." + ) from error + object.__setattr__( + self, + "batch_item_id", + _normalize_int(self.batch_item_id, name="batch_item_id"), + ) + + +@dataclass(frozen=True) +class SidePackingSpec: + """Controller input describing one model side's exact token geometry. + + Length tuples align one-to-one with the canonical ``items`` passed to + :func:`build_lockstep_packing_plan`. ``raw_lengths`` are exact tokenizer + lengths before backend padding. ``effective_lengths`` include every + per-sample padding rule. ``physical_size_fn``, when supplied, receives the + ordered effective lengths in a candidate bin and returns its complete + physical size, including packed-tail padding. It must be deterministic and + may not return a value smaller than the sum of the effective lengths. + """ + + side_id: str + capacity: int + raw_lengths: tuple[int, ...] + effective_lengths: tuple[int, ...] + physical_size_fn: PhysicalSizeFn | None = None + + def __post_init__(self) -> None: + if not isinstance(self.side_id, str) or not self.side_id: + raise ValueError( + f"side_id must be a non-empty string, got {self.side_id!r}." + ) + object.__setattr__( + self, "capacity", _normalize_int(self.capacity, name="capacity", minimum=1) + ) + raw_lengths = _normalize_int_tuple( + self.raw_lengths, name=f"{self.side_id}.raw_lengths", minimum=1 + ) + effective_lengths = _normalize_int_tuple( + self.effective_lengths, + name=f"{self.side_id}.effective_lengths", + minimum=1, + ) + if len(raw_lengths) != len(effective_lengths): + raise ValueError( + f"side {self.side_id!r} has {len(raw_lengths)} raw lengths but " + f"{len(effective_lengths)} effective lengths." + ) + for ordinal, (raw, effective) in enumerate( + zip(raw_lengths, effective_lengths, strict=True) + ): + if effective < raw: + raise ValueError( + f"side {self.side_id!r} item {ordinal} has effective length " + f"{effective} smaller than raw length {raw}." + ) + if self.physical_size_fn is not None and not callable(self.physical_size_fn): + raise ValueError( + f"side {self.side_id!r} physical_size_fn must be callable or None." + ) + object.__setattr__(self, "raw_lengths", raw_lengths) + object.__setattr__(self, "effective_lengths", effective_lengths) + + def physical_size(self, item_ordinals: Sequence[int]) -> int: + """Return the complete physical size of an ordered candidate bin.""" + effective_lengths = tuple(self.effective_lengths[i] for i in item_ordinals) + if not effective_lengths: + raise ValueError(f"side {self.side_id!r} cannot size an empty bin.") + minimum_size = sum(effective_lengths) + if self.physical_size_fn is None: + return minimum_size + try: + physical_size = self.physical_size_fn(effective_lengths) + except Exception as error: + raise ValueError( + f"side {self.side_id!r} physical_size_fn failed for effective " + f"lengths {effective_lengths}." + ) from error + physical_size = _normalize_int( + physical_size, + name=f"side {self.side_id!r} physical size", + minimum=1, + ) + if physical_size < minimum_size: + raise ValueError( + f"side {self.side_id!r} physical_size_fn returned {physical_size} " + f"for effective lengths {effective_lengths}, smaller than their " + f"sum {minimum_size}." + ) + return physical_size + + +@dataclass(frozen=True) +class SidePackingPlan: + """Immutable model-specific geometry for shared logical bins.""" + + side_id: str + capacity: int + raw_lengths: tuple[int, ...] + effective_lengths: tuple[int, ...] + raw_cu_seqlens_by_bin: tuple[tuple[int, ...], ...] + padded_cu_seqlens_by_bin: tuple[tuple[int, ...], ...] + physical_tokens_by_bin: tuple[int, ...] + rank_bin_indices: tuple[tuple[int, ...], ...] + + def __post_init__(self) -> None: + if not isinstance(self.side_id, str) or not self.side_id: + raise ValueError( + f"side_id must be a non-empty string, got {self.side_id!r}." + ) + object.__setattr__( + self, "capacity", _normalize_int(self.capacity, name="capacity", minimum=1) + ) + object.__setattr__( + self, + "raw_lengths", + _normalize_int_tuple( + self.raw_lengths, name=f"{self.side_id}.raw_lengths", minimum=1 + ), + ) + object.__setattr__( + self, + "effective_lengths", + _normalize_int_tuple( + self.effective_lengths, + name=f"{self.side_id}.effective_lengths", + minimum=1, + ), + ) + if len(self.raw_lengths) != len(self.effective_lengths): + raise ValueError( + f"side {self.side_id!r} raw/effective length counts do not match." + ) + for ordinal, (raw, effective) in enumerate( + zip(self.raw_lengths, self.effective_lengths, strict=True) + ): + if effective < raw: + raise ValueError( + f"side {self.side_id!r} item {ordinal} has effective length " + f"{effective} smaller than raw length {raw}." + ) + + raw_cu = tuple( + _normalize_int_tuple( + values, + name=f"{self.side_id}.raw_cu_seqlens_by_bin[{bin_idx}]", + ) + for bin_idx, values in enumerate(self.raw_cu_seqlens_by_bin) + ) + padded_cu = tuple( + _normalize_int_tuple( + values, + name=f"{self.side_id}.padded_cu_seqlens_by_bin[{bin_idx}]", + ) + for bin_idx, values in enumerate(self.padded_cu_seqlens_by_bin) + ) + physical_tokens = _normalize_int_tuple( + self.physical_tokens_by_bin, + name=f"{self.side_id}.physical_tokens_by_bin", + minimum=1, + ) + rank_bin_indices = tuple( + _normalize_int_tuple( + values, name=f"{self.side_id}.rank_bin_indices[{rank_idx}]" + ) + for rank_idx, values in enumerate(self.rank_bin_indices) + ) + object.__setattr__(self, "raw_cu_seqlens_by_bin", raw_cu) + object.__setattr__(self, "padded_cu_seqlens_by_bin", padded_cu) + object.__setattr__(self, "physical_tokens_by_bin", physical_tokens) + object.__setattr__(self, "rank_bin_indices", rank_bin_indices) + + bin_count = len(physical_tokens) + if bin_count == 0: + raise ValueError(f"side {self.side_id!r} must contain at least one bin.") + if len(raw_cu) != bin_count or len(padded_cu) != bin_count: + raise ValueError( + f"side {self.side_id!r} bin metadata counts differ: raw_cu=" + f"{len(raw_cu)}, padded_cu={len(padded_cu)}, physical={bin_count}." + ) + for bin_idx, (raw, padded, physical) in enumerate( + zip(raw_cu, padded_cu, physical_tokens, strict=True) + ): + if len(raw) < 2 or raw[0] != 0: + raise ValueError( + f"side {self.side_id!r} bin {bin_idx} raw cu_seqlens must " + "start at zero and describe at least one item." + ) + if len(padded) != len(raw) or padded[0] != 0: + raise ValueError( + f"side {self.side_id!r} bin {bin_idx} raw and padded " + "cu_seqlens must have the same shape and start at zero." + ) + if any(left >= right for left, right in zip(raw, raw[1:])): + raise ValueError( + f"side {self.side_id!r} bin {bin_idx} raw cu_seqlens must " + "be strictly increasing." + ) + if any(left >= right for left, right in zip(padded, padded[1:])): + raise ValueError( + f"side {self.side_id!r} bin {bin_idx} padded cu_seqlens " + "must be strictly increasing." + ) + if padded[-1] != physical: + raise ValueError( + f"side {self.side_id!r} bin {bin_idx} padded cu_seqlens end " + f"at {padded[-1]}, not physical size {physical}." + ) + if raw[-1] > physical: + raise ValueError( + f"side {self.side_id!r} bin {bin_idx} raw size {raw[-1]} " + f"exceeds physical size {physical}." + ) + if physical > self.capacity: + raise ValueError( + f"side {self.side_id!r} bin {bin_idx} uses {physical} " + f"physical tokens, exceeding capacity {self.capacity}." + ) + + if not rank_bin_indices or any(not indices for indices in rank_bin_indices): + raise ValueError( + f"side {self.side_id!r} must assign at least one bin to every DP rank." + ) + flattened_rank_bins = tuple( + bin_idx for indices in rank_bin_indices for bin_idx in indices + ) + if flattened_rank_bins != tuple(range(bin_count)): + raise ValueError( + f"side {self.side_id!r} rank_bin_indices must cover bins once in " + f"contiguous rank order; got {rank_bin_indices}." + ) + rank_bin_counts = {len(indices) for indices in rank_bin_indices} + if len(rank_bin_counts) != 1: + raise ValueError( + f"side {self.side_id!r} has unequal DP bin counts " + f"{tuple(len(indices) for indices in rank_bin_indices)}." + ) + + +class _FrozenSidePlans(Mapping[str, SidePackingPlan]): + """Small immutable and pickle-friendly mapping used by a lockstep plan.""" + + __slots__ = ("__items",) + + def __init__(self, values: Mapping[str, SidePackingPlan]): + object.__setattr__(self, "_FrozenSidePlans__items", tuple(values.items())) + + def __getitem__(self, key: str) -> SidePackingPlan: + for candidate, value in self.__items: + if candidate == key: + return value + raise KeyError(key) + + def __iter__(self) -> Iterator[str]: + return (key for key, _ in self.__items) + + def __len__(self) -> int: + return len(self.__items) + + def __hash__(self) -> int: + return hash(frozenset(self.__items)) + + def __repr__(self) -> str: + return repr(dict(self.__items)) + + def __reduce__(self) -> tuple[object, tuple[dict[str, SidePackingPlan]]]: + return type(self), (dict(self.__items),) + + def __setattr__(self, name: str, value: object) -> None: + raise AttributeError(f"{type(self).__name__} is immutable.") + + def __delattr__(self, name: str) -> None: + raise AttributeError(f"{type(self).__name__} is immutable.") + + +@dataclass(frozen=True) +class LockstepPackingPlan: + """Immutable shared membership and side-local layout for one global batch.""" + + batch_uid: int + canonical_batch_item_ids: tuple[int, ...] + bins: tuple[tuple[int, ...], ...] + sides: Mapping[str, SidePackingPlan] + equalization_splits: int = 0 + + def __post_init__(self) -> None: + object.__setattr__( + self, "batch_uid", _normalize_int(self.batch_uid, name="batch_uid") + ) + object.__setattr__( + self, + "equalization_splits", + _normalize_int(self.equalization_splits, name="equalization_splits"), + ) + canonical_ids = _normalize_int_tuple( + self.canonical_batch_item_ids, name="canonical_batch_item_ids" + ) + bins = tuple( + _normalize_int_tuple(values, name=f"bins[{bin_idx}]") + for bin_idx, values in enumerate(self.bins) + ) + object.__setattr__(self, "canonical_batch_item_ids", canonical_ids) + object.__setattr__(self, "bins", bins) + + if not canonical_ids: + raise ValueError("LockstepPackingPlan requires at least one batch item.") + if len(set(canonical_ids)) != len(canonical_ids): + raise ValueError( + "canonical_batch_item_ids must be occurrence-unique; got " + f"{canonical_ids}." + ) + if not bins or any(not values for values in bins): + raise ValueError("LockstepPackingPlan bins must all be non-empty.") + flattened_ids = tuple( + batch_item_id for values in bins for batch_item_id in values + ) + if len(flattened_ids) != len(canonical_ids) or set(flattened_ids) != set( + canonical_ids + ): + raise ValueError( + "LockstepPackingPlan bins must contain every canonical " + "batch_item_id exactly once." + ) + + if not isinstance(self.sides, Mapping) or not self.sides: + raise ValueError("LockstepPackingPlan requires at least one side.") + frozen_sides = _FrozenSidePlans(dict(self.sides)) + object.__setattr__(self, "sides", frozen_sides) + + id_to_ordinal = { + batch_item_id: ordinal + for ordinal, batch_item_id in enumerate(canonical_ids) + } + reference_routing: tuple[tuple[int, ...], ...] | None = None + for side_key, side in frozen_sides.items(): + if not isinstance(side, SidePackingPlan): + raise ValueError( + f"side {side_key!r} must be a SidePackingPlan, got " + f"{type(side).__name__}." + ) + if side_key != side.side_id: + raise ValueError( + f"side mapping key {side_key!r} does not match plan side_id " + f"{side.side_id!r}." + ) + if len(side.raw_lengths) != len(canonical_ids): + raise ValueError( + f"side {side.side_id!r} has {len(side.raw_lengths)} lengths " + f"for {len(canonical_ids)} canonical batch items." + ) + if len(side.physical_tokens_by_bin) != len(bins): + raise ValueError( + f"side {side.side_id!r} has geometry for " + f"{len(side.physical_tokens_by_bin)} bins, expected {len(bins)}." + ) + if reference_routing is None: + reference_routing = side.rank_bin_indices + elif side.rank_bin_indices != reference_routing: + raise ValueError( + f"side {side.side_id!r} rank routing differs from the shared " + "lockstep routing." + ) + + for bin_idx, batch_item_ids in enumerate(bins): + ordinals = tuple(id_to_ordinal[item_id] for item_id in batch_item_ids) + raw_lengths = tuple(side.raw_lengths[i] for i in ordinals) + effective_lengths = tuple(side.effective_lengths[i] for i in ordinals) + expected_raw_cu = _cumulative_lengths(raw_lengths) + if side.raw_cu_seqlens_by_bin[bin_idx] != expected_raw_cu: + raise ValueError( + f"side {side.side_id!r} bin {bin_idx} raw cu_seqlens do " + "not match shared bin membership." + ) + expected_effective_cu = _cumulative_lengths(effective_lengths) + padded_cu = side.padded_cu_seqlens_by_bin[bin_idx] + if padded_cu[:-1] != expected_effective_cu[:-1]: + raise ValueError( + f"side {side.side_id!r} bin {bin_idx} padded cu_seqlens " + "do not preserve per-item effective boundaries." + ) + if padded_cu[-1] < expected_effective_cu[-1]: + raise ValueError( + f"side {side.side_id!r} bin {bin_idx} padded size " + f"{padded_cu[-1]} is smaller than effective size " + f"{expected_effective_cu[-1]}." + ) + + assert reference_routing is not None + dp_size = len(reference_routing) + if len(canonical_ids) % dp_size != 0: + raise ValueError( + f"{len(canonical_ids)} canonical items cannot be divided into " + f"{dp_size} equal contiguous DP shards." + ) + items_per_rank = len(canonical_ids) // dp_size + for rank, rank_bins in enumerate(reference_routing): + rank_ordinals = { + id_to_ordinal[item_id] + for bin_idx in rank_bins + for item_id in bins[bin_idx] + } + expected_ordinals = set( + range(rank * items_per_rank, (rank + 1) * items_per_rank) + ) + if rank_ordinals != expected_ordinals: + raise ValueError( + f"DP rank {rank} bins do not contain exactly its canonical " + "contiguous logical shard." + ) + + +def _pack_dp_shard( + item_ordinals: Sequence[int], side_specs: Sequence[SidePackingSpec] +) -> list[list[int]]: + pressures = { + item_ordinal: max( + Fraction(side.effective_lengths[item_ordinal], side.capacity) + for side in side_specs + ) + for item_ordinal in item_ordinals + } + decreasing = sorted( + item_ordinals, key=lambda item_ordinal: (-pressures[item_ordinal], item_ordinal) + ) + + bins: list[list[int]] = [] + for item_ordinal in decreasing: + # First-fit is lockstep across model sides: reuse the first existing bin + # only when the candidate's complete side-local geometry fits every side. + for bin_items in bins: + candidate = (*bin_items, item_ordinal) + if all( + side.physical_size(candidate) <= side.capacity for side in side_specs + ): + bin_items.append(item_ordinal) + break + else: + # Every item was prevalidated to fit an empty bin. If no existing + # bin has room on all sides, retain the item as the new bin's first. + bins.append([item_ordinal]) + return bins + + +def _equalize_rank_bin_counts( + rank_bins: list[list[list[int]]], *, canonical_item_count: int +) -> int: + target_count = max(len(bins) for bins in rank_bins) + split_count = 0 + for rank, bins in enumerate(rank_bins): + while len(bins) < target_count: + splittable = [ + bin_idx for bin_idx, bin_items in enumerate(bins) if len(bin_items) > 1 + ] + if not splittable: + raise ValueError( + f"DP rank {rank} has {len(bins)} bins and cannot reach the " + f"required equalized count {target_count} by splitting its " + f"logical samples (global batch has {canonical_item_count} items)." + ) + source_bin_idx = max(splittable, key=lambda i: (len(bins[i]), -i)) + moved_item = bins[source_bin_idx].pop() + bins.append([moved_item]) + split_count += 1 + return split_count + + +def build_lockstep_packing_plan( + *, + batch_uid: int, + items: Sequence[LockstepPackingItem], + sides: Sequence[SidePackingSpec], + data_parallel_size: int, +) -> LockstepPackingPlan: + """Build one deterministic multi-side FFD plan over contiguous DP shards. + + The input order is the canonical global-batch order. It is divided evenly + into contiguous DP shards before packing. Each shard is packed independently + with stable multi-dimensional first-fit-decreasing. Ranks with fewer bins are + equalized only by deterministic subset splitting; logical ownership never + changes. + """ + batch_uid = _normalize_int(batch_uid, name="batch_uid") + data_parallel_size = _normalize_int( + data_parallel_size, name="data_parallel_size", minimum=1 + ) + canonical_items = tuple(items) + side_specs = tuple(sides) + if not canonical_items: + raise ValueError("Cannot build a lockstep packing plan for an empty batch.") + if not side_specs: + raise ValueError("Cannot build a lockstep packing plan without model sides.") + if any(not isinstance(item, LockstepPackingItem) for item in canonical_items): + raise ValueError("items must contain only LockstepPackingItem records.") + if any(not isinstance(side, SidePackingSpec) for side in side_specs): + raise ValueError("sides must contain only SidePackingSpec records.") + + batch_item_ids = tuple(item.batch_item_id for item in canonical_items) + if len(set(batch_item_ids)) != len(batch_item_ids): + duplicates = sorted( + item_id + for item_id in set(batch_item_ids) + if batch_item_ids.count(item_id) > 1 + ) + raise ValueError( + "batch_item_id must be unique for every logical occurrence; duplicate " + f"IDs: {duplicates}." + ) + side_ids = tuple(side.side_id for side in side_specs) + if len(set(side_ids)) != len(side_ids): + raise ValueError(f"side_id values must be unique, got {side_ids}.") + if len(canonical_items) % data_parallel_size != 0: + raise ValueError( + f"global batch has {len(canonical_items)} logical items, which is not " + f"divisible by data_parallel_size={data_parallel_size}. Pad validation " + "batches before assigning batch_item_id and planning." + ) + for side in side_specs: + if len(side.raw_lengths) != len(canonical_items): + raise ValueError( + f"side {side.side_id!r} has {len(side.raw_lengths)} lengths for " + f"{len(canonical_items)} logical items." + ) + + for item_ordinal, item in enumerate(canonical_items): + for side in side_specs: + physical_tokens = side.physical_size((item_ordinal,)) + if physical_tokens > side.capacity: + raise ValueError( + f"sample_id={item.sample_id!r} (batch_item_id={item.batch_item_id}) " + f"does not fit side {side.side_id!r}: raw_length=" + f"{side.raw_lengths[item_ordinal]}, effective_length=" + f"{side.effective_lengths[item_ordinal]}, physical_tokens=" + f"{physical_tokens}, capacity={side.capacity}." + ) + + items_per_rank = len(canonical_items) // data_parallel_size + rank_bins = [ + _pack_dp_shard( + range(rank * items_per_rank, (rank + 1) * items_per_rank), side_specs + ) + for rank in range(data_parallel_size) + ] + equalization_splits = _equalize_rank_bin_counts( + rank_bins, canonical_item_count=len(canonical_items) + ) + + flattened_ordinal_bins: list[list[int]] = [] + rank_bin_indices: list[tuple[int, ...]] = [] + for bins in rank_bins: + first_bin_idx = len(flattened_ordinal_bins) + flattened_ordinal_bins.extend(bins) + rank_bin_indices.append(tuple(range(first_bin_idx, first_bin_idx + len(bins)))) + shared_bins = tuple( + tuple(batch_item_ids[item_ordinal] for item_ordinal in bin_items) + for bin_items in flattened_ordinal_bins + ) + shared_rank_bin_indices = tuple(rank_bin_indices) + + side_plans: dict[str, SidePackingPlan] = {} + for side in side_specs: + raw_cu_by_bin: list[tuple[int, ...]] = [] + padded_cu_by_bin: list[tuple[int, ...]] = [] + physical_tokens_by_bin: list[int] = [] + for bin_items in flattened_ordinal_bins: + raw_lengths = tuple(side.raw_lengths[i] for i in bin_items) + effective_lengths = tuple(side.effective_lengths[i] for i in bin_items) + physical_tokens = side.physical_size(bin_items) + if physical_tokens > side.capacity: + raise ValueError( + f"internal lockstep packing error: side {side.side_id!r} bin " + f"{tuple(batch_item_ids[i] for i in bin_items)} uses " + f"{physical_tokens} physical tokens, exceeding capacity " + f"{side.capacity}." + ) + raw_cu = _cumulative_lengths(raw_lengths) + padded_cu = list(_cumulative_lengths(effective_lengths)) + padded_cu[-1] = physical_tokens + raw_cu_by_bin.append(raw_cu) + padded_cu_by_bin.append(tuple(padded_cu)) + physical_tokens_by_bin.append(physical_tokens) + + side_plans[side.side_id] = SidePackingPlan( + side_id=side.side_id, + capacity=side.capacity, + raw_lengths=side.raw_lengths, + effective_lengths=side.effective_lengths, + raw_cu_seqlens_by_bin=tuple(raw_cu_by_bin), + padded_cu_seqlens_by_bin=tuple(padded_cu_by_bin), + physical_tokens_by_bin=tuple(physical_tokens_by_bin), + rank_bin_indices=shared_rank_bin_indices, + ) + + return LockstepPackingPlan( + batch_uid=batch_uid, + canonical_batch_item_ids=batch_item_ids, + bins=shared_bins, + sides=side_plans, + equalization_splits=equalization_splits, + ) diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index 0d1d4e17e66..935fba91ac6 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -992,6 +992,8 @@ def kd_data_processor( "kd_data_processor expected 'tools' to be a list of objects or None" ) output["tools"] = deepcopy(tools) + if "sample_id" in datum_dict: + output["sample_id"] = datum_dict["sample_id"] return output diff --git a/nemo_rl/distributed/batched_data_dict.py b/nemo_rl/distributed/batched_data_dict.py index ccad97a80f0..ddfbc4e475e 100644 --- a/nemo_rl/distributed/batched_data_dict.py +++ b/nemo_rl/distributed/batched_data_dict.py @@ -38,7 +38,7 @@ VLLM_PROMPT_KEYS, PackedTensor, ) -from nemo_rl.data.packing import get_packer +from nemo_rl.data.packing import LockstepPackingPlan, get_packer from nemo_rl.data_plane.schema import OPD_FULL_FIELDS from nemo_rl.distributed.collectives import ( gather_jagged_object_lists, @@ -152,6 +152,9 @@ def __init__(self, *args, **kwargs): self.micro_batch_indices = None self.micro_batch_lengths = None self.elem_counts_per_gb = None + self.lockstep_packing_plan: LockstepPackingPlan | None = None + self.lockstep_side_id: str | None = None + self.lockstep_bin_indices: tuple[int, ...] | None = None def get_multimodal_dict( self, @@ -955,6 +958,89 @@ def _get_padded_seqlen(seqlen: int) -> int: return aggregated_shards + def shard_by_packing_plan( + self, + plan: LockstepPackingPlan, + *, + side_id: str, + ) -> list["SlicedDataDict"]: + """Shard one logical global batch using a controller-owned lockstep plan. + + Unlike :meth:`shard_by_batch_size`, this method never computes or + repairs bin membership. It validates occurrence IDs, follows the + plan's canonical contiguous DP ownership, and installs the physical-bin + schedule expected by backend microbatch iterators. + """ + if side_id not in plan.sides: + raise ValueError( + f"lockstep plan batch_uid={plan.batch_uid} has no side " + f"{side_id!r}; available sides are {tuple(plan.sides)!r}." + ) + if "batch_item_id" not in self.data: + raise ValueError( + "plan-driven sharding requires a batch_item_id field assigned " + "by the CSD controller." + ) + item_ids_value = self.data["batch_item_id"] + if isinstance(item_ids_value, torch.Tensor): + if item_ids_value.ndim != 1: + raise ValueError( + "batch_item_id must be a one-dimensional tensor or list; " + f"got shape {tuple(item_ids_value.shape)}." + ) + item_ids = tuple(int(value) for value in item_ids_value.tolist()) + elif isinstance(item_ids_value, list): + item_ids = tuple(int(value) for value in item_ids_value) + else: + raise TypeError( + "batch_item_id must be a one-dimensional tensor or list, got " + f"{type(item_ids_value).__name__}." + ) + if len(set(item_ids)) != len(item_ids): + raise ValueError( + "batch_item_id values must be occurrence-unique before " + f"plan-driven sharding; got {item_ids}." + ) + if len(item_ids) != len(plan.canonical_batch_item_ids) or set(item_ids) != set( + plan.canonical_batch_item_ids + ): + missing = sorted(set(plan.canonical_batch_item_ids) - set(item_ids)) + unexpected = sorted(set(item_ids) - set(plan.canonical_batch_item_ids)) + raise ValueError( + f"batch_item_id membership differs from lockstep plan " + f"batch_uid={plan.batch_uid}: missing={missing}, " + f"unexpected={unexpected}." + ) + + row_by_id = {item_id: row for row, item_id in enumerate(item_ids)} + side_plan = plan.sides[side_id] + shards: list[SlicedDataDict] = [] + for rank_bin_indices in side_plan.rank_bin_indices: + rank_bins = [plan.bins[bin_index] for bin_index in rank_bin_indices] + rank_item_ids = [item_id for values in rank_bins for item_id in values] + rank_rows = [row_by_id[item_id] for item_id in rank_item_ids] + shard = SlicedDataDict(self.select_indices(rank_rows).data) + + micro_batch_indices: list[list[int]] = [] + row_offset = 0 + for values in rank_bins: + next_offset = row_offset + len(values) + micro_batch_indices.append([row_offset, next_offset]) + row_offset = next_offset + shard.micro_batch_indices = [micro_batch_indices] + shard.micro_batch_lengths = [ + [ + side_plan.physical_tokens_by_bin[bin_index] + for bin_index in rank_bin_indices + ] + ] + shard.elem_counts_per_gb = [len(rank_item_ids)] + shard.lockstep_packing_plan = plan + shard.lockstep_side_id = side_id + shard.lockstep_bin_indices = rank_bin_indices + shards.append(shard) + return shards + def get_batch(self, batch_idx, batch_size=None) -> "SlicedDataDict": """Slices a subbatch from the batch. @@ -1015,6 +1101,13 @@ def slice(self, start: int, end: int) -> "SlicedDataDict": if isinstance(v, torch.Tensor) and v.is_nested else v[start:end] ) + # A lockstep plan is immutable controller metadata, not a batch-axis + # payload. Preserve it through the global-batch and physical-bin slices + # taken inside workers so each backend can validate the materialized + # IDs and geometry instead of silently recomputing membership. + sliced_batch.lockstep_packing_plan = self.lockstep_packing_plan + sliced_batch.lockstep_side_id = self.lockstep_side_id + sliced_batch.lockstep_bin_indices = self.lockstep_bin_indices return sliced_batch def repeat_interleave( diff --git a/nemo_rl/distributed/model_utils.py b/nemo_rl/distributed/model_utils.py index 18df79d8e1c..f753b0dcd3a 100644 --- a/nemo_rl/distributed/model_utils.py +++ b/nemo_rl/distributed/model_utils.py @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import math from collections.abc import Callable from typing import TYPE_CHECKING, Any, Optional import torch import torch.distributed.nn.functional -from torch.distributed.tensor import DTensor +from torch.distributed.tensor import DTensor, distribute_tensor from nemo_rl.algorithms.logits_sampling_utils import ( TrainingSamplingParams, @@ -225,6 +226,272 @@ def backward( return grad_input, None, None, None, None, None, None +class DistributedSelectedLogprobs(torch.autograd.Function): + """Exact selected-token log-probs without gathering a TP vocab shard. + + ``batch_indices`` / ``seq_indices`` select ``N`` local (possibly CP + sharded) rows and ``token_ids`` supplies ``K`` global vocabulary ids per + row. Forward communicates only row-wise normalization scalars and the + selected values. Backward rematerializes the familiar + ``one_hot(target)-softmax`` gradient on each local vocabulary shard, summed + over the selected supports, and scatters it back into the original rows. + + Columns at or beyond ``real_vocab_size`` are excluded from both logZ and + gradients. This matters for TP-divisibility padding (for example a + Megatron LM head with 248320 columns for a 248077-token tokenizer). + """ + + @staticmethod + def forward( # pyrefly: ignore[bad-override] + ctx: Any, + vocab_parallel_logits: torch.Tensor, + batch_indices: torch.Tensor, + seq_indices: torch.Tensor, + token_ids: torch.Tensor, + vocab_start_index: int, + vocab_end_index: int, + real_vocab_size: int, + temperature: float, + group: Optional[torch.distributed.ProcessGroup], + ) -> torch.Tensor: + if temperature <= 0: + raise ValueError(f"temperature must be > 0, got {temperature}") + if batch_indices.dim() != 1 or seq_indices.dim() != 1: + raise ValueError("batch_indices and seq_indices must be one-dimensional") + if batch_indices.shape != seq_indices.shape: + raise ValueError( + "batch_indices and seq_indices must have the same shape; got " + f"{tuple(batch_indices.shape)} and {tuple(seq_indices.shape)}" + ) + squeeze = token_ids.dim() == 1 + if squeeze: + token_ids = token_ids.unsqueeze(-1) + if token_ids.dim() != 2 or token_ids.shape[0] != batch_indices.shape[0]: + raise ValueError( + "token_ids must have shape [N] or [N, K] matching the selected " + f"rows; got {tuple(token_ids.shape)} for N={batch_indices.shape[0]}" + ) + + rows = vocab_parallel_logits[batch_indices, seq_indices].float() + rows.div_(float(temperature)) + local_vocab = rows.shape[-1] + valid_local = max( + 0, + min(int(vocab_end_index), int(real_vocab_size)) - int(vocab_start_index), + ) + valid_local = min(valid_local, local_vocab) + if valid_local < local_vocab: + rows[:, valid_local:] = -torch.inf + + tp_world = ( + torch.distributed.get_world_size(group) + if group is not None and torch.distributed.is_initialized() + else 1 + ) + + target_valid = (token_ids >= 0) & (token_ids < int(real_vocab_size)) + target_owned = ( + target_valid + & (token_ids >= int(vocab_start_index)) + & (token_ids < int(vocab_start_index) + valid_local) + ) + local_target = ( + (token_ids - int(vocab_start_index)) + .clamp(min=0, max=max(local_vocab - 1, 0)) + .long() + ) + selected = torch.gather(rows, dim=-1, index=local_target) + selected = torch.where(target_owned, selected, torch.zeros_like(selected)) + + row_max = rows.max(dim=-1, keepdim=True).values + if tp_world > 1: + torch.distributed.all_reduce( + row_max, op=torch.distributed.ReduceOp.MAX, group=group + ) + rows.sub_(row_max).exp_() + row_sum = rows.sum(dim=-1, keepdim=True) + if tp_world > 1: + torch.distributed.all_reduce( + row_sum, op=torch.distributed.ReduceOp.SUM, group=group + ) + log_z = row_max + row_sum.log() + + if tp_world > 1: + torch.distributed.all_reduce( + selected, op=torch.distributed.ReduceOp.SUM, group=group + ) + selected = selected - log_z + selected = torch.where( + target_valid, selected, torch.full_like(selected, -torch.inf) + ) + + ctx.save_for_backward( + vocab_parallel_logits, + target_owned, + local_target, + batch_indices, + seq_indices, + ) + ctx.input_shape = tuple(vocab_parallel_logits.shape) + ctx.input_dtype = vocab_parallel_logits.dtype + ctx.temperature = float(temperature) + ctx.valid_local = int(valid_local) + ctx.group = group + ctx.tp_world = int(tp_world) + ctx.squeeze = squeeze + return selected.squeeze(-1) if squeeze else selected + + @staticmethod + def backward( + ctx: Any, + grad_output: torch.Tensor, + ) -> tuple[ + torch.Tensor, + None, + None, + None, + None, + None, + None, + None, + None, + ]: + logits, target_owned, local_target, batch_indices, seq_indices = ( + ctx.saved_tensors + ) + if ctx.squeeze: + grad_output = grad_output.unsqueeze(-1) + grad_output = grad_output.float() + + # Recompute the local softmax shard instead of saving an fp32 + # ``[N, V_local]`` tensor from forward. Saving ``logits`` above retains + # only a reference to the model output that autograd already owns; + # these temporaries are released as soon as this backward returns. + rows = logits[batch_indices, seq_indices].float() + rows.div_(ctx.temperature) + local_vocab = int(rows.shape[-1]) + if ctx.valid_local < local_vocab: + rows[:, ctx.valid_local :] = -torch.inf + row_max = rows.max(dim=-1, keepdim=True).values + if ctx.tp_world > 1: + torch.distributed.all_reduce( + row_max, + op=torch.distributed.ReduceOp.MAX, + group=ctx.group, + ) + rows.sub_(row_max).exp_() + row_sum = rows.sum(dim=-1, keepdim=True) + if ctx.tp_world > 1: + torch.distributed.all_reduce( + row_sum, + op=torch.distributed.ReduceOp.SUM, + group=ctx.group, + ) + softmax = rows.div_(row_sum) + grad_rows = -softmax * grad_output.sum(dim=-1, keepdim=True) + + n_rows = int(softmax.shape[0]) + if n_rows: + row_offsets = ( + torch.arange(n_rows, device=softmax.device, dtype=torch.long) + * local_vocab + ).unsqueeze(-1) + flat_selected = row_offsets + local_target + owned_grad = torch.where( + target_owned, grad_output, torch.zeros_like(grad_output) + ) + grad_rows.view(-1).scatter_add_( + 0, flat_selected.reshape(-1), owned_grad.reshape(-1) + ) + grad_rows.div_(ctx.temperature) + + grad_input = torch.zeros( + ctx.input_shape, device=softmax.device, dtype=ctx.input_dtype + ) + flat_rows = batch_indices * ctx.input_shape[1] + seq_indices + grad_input.view(-1, ctx.input_shape[-1]).index_add_( + 0, flat_rows, grad_rows.to(ctx.input_dtype) + ) + return ( + grad_input, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +def distributed_selected_logprobs( + vocab_parallel_logits: torch.Tensor, + batch_indices: torch.Tensor, + seq_indices: torch.Tensor, + token_ids: torch.Tensor, + *, + vocab_start_index: int, + vocab_end_index: int, + real_vocab_size: int, + temperature: float, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + row_chunk_size: Optional[int] = None, +) -> torch.Tensor: + """Selected-token log-probs over a TP-sharded vocabulary. + + Functional wrapper around :class:`DistributedSelectedLogprobs` that tiles + the selected rows into chunks of ``row_chunk_size`` to bound the fp32 + ``[rows, V_local]`` working set. + + Args: + vocab_parallel_logits: Local logits ``[B, S_local, V_local]``. + batch_indices: Row batch indices ``[N]``. + seq_indices: Row sequence indices ``[N]`` into ``S_local``. + token_ids: Global vocabulary ids ``[N]`` or ``[N, K]``; ids outside + ``[0, real_vocab_size)`` return ``-inf``. + vocab_start_index: First global vocabulary id held by this TP rank. + vocab_end_index: One past the last global vocabulary id held here. + real_vocab_size: Tokenizer vocabulary size; padded columns are ignored. + temperature: Softmax temperature applied to the logits. + tp_group: Tensor-parallel group, or None for an unsharded vocabulary. + row_chunk_size: Rows per tile; None or <= 0 evaluates all rows at once. + + Returns: + Log-probabilities with the shape of ``token_ids``. + """ + n_rows = int(batch_indices.shape[0]) + if row_chunk_size is None or row_chunk_size <= 0 or n_rows <= row_chunk_size: + return DistributedSelectedLogprobs.apply( + vocab_parallel_logits, + batch_indices, + seq_indices, + token_ids, + vocab_start_index, + vocab_end_index, + real_vocab_size, + temperature, + tp_group, + ) + chunks: list[torch.Tensor] = [] + for start in range(0, n_rows, int(row_chunk_size)): + end = min(start + int(row_chunk_size), n_rows) + chunks.append( + DistributedSelectedLogprobs.apply( + vocab_parallel_logits, + batch_indices[start:end], + seq_indices[start:end], + token_ids[start:end], + vocab_start_index, + vocab_end_index, + real_vocab_size, + temperature, + tp_group, + ) + ) + return torch.cat(chunks, dim=0) + + class ChunkedDistributedLogprob(torch.autograd.Function): """Custom autograd function for computing log probabilities in a distributed setting. @@ -338,6 +605,107 @@ def backward( return grad_input, None, None, None, None, None, None +class ChunkedVocabParallelLogsumexp(torch.autograd.Function): + """Memory-bounded full-vocabulary logsumexp for TP-sharded logits. + + Forward computes one sequence chunk at a time and saves only the original + logits. Backward recomputes the distributed softmax chunk-by-chunk, avoiding + a persistent full-sequence fp32 softmax while preserving exact gradients on + each TP vocabulary shard. + """ + + @staticmethod + def forward( # pyrefly: ignore[bad-override] + ctx: Any, + vocab_parallel_logits: torch.Tensor, + temperature: float, + chunk_size: int, + tp_group: torch.distributed.ProcessGroup, + ) -> torch.Tensor: + work_dtype = ( + torch.float32 + if vocab_parallel_logits.dtype in (torch.float16, torch.bfloat16) + else vocab_parallel_logits.dtype + ) + batch_size, seq_size, _ = vocab_parallel_logits.shape + output = torch.empty( + (batch_size, seq_size), + dtype=work_dtype, + device=vocab_parallel_logits.device, + ) + + for chunk_start in range(0, seq_size, chunk_size): + chunk_end = min(chunk_start + chunk_size, seq_size) + scaled = ( + vocab_parallel_logits[:, chunk_start:chunk_end, :].to(work_dtype) + / temperature + ) + logits_max = torch.amax(scaled, dim=-1) + torch.distributed.all_reduce( + logits_max, + op=torch.distributed.ReduceOp.MAX, + group=tp_group, + ) + scaled.sub_(logits_max.unsqueeze(-1)).exp_() + sum_exp = scaled.sum(dim=-1) + torch.distributed.all_reduce( + sum_exp, + op=torch.distributed.ReduceOp.SUM, + group=tp_group, + ) + output[:, chunk_start:chunk_end] = sum_exp.log_().add_(logits_max) + + ctx.save_for_backward(vocab_parallel_logits) + ctx.temperature = temperature + ctx.chunk_size = chunk_size + ctx.tp_group = tp_group + return output + + @staticmethod + def backward( # pyrefly: ignore[bad-override] + ctx: Any, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, None, None, None]: + (vocab_parallel_logits,) = ctx.saved_tensors + temperature = ctx.temperature + chunk_size = ctx.chunk_size + tp_group = ctx.tp_group + work_dtype = ( + torch.float32 + if vocab_parallel_logits.dtype in (torch.float16, torch.bfloat16) + else vocab_parallel_logits.dtype + ) + seq_size = int(vocab_parallel_logits.shape[1]) + grad_input = torch.empty_like(vocab_parallel_logits) + + for chunk_start in range(0, seq_size, chunk_size): + chunk_end = min(chunk_start + chunk_size, seq_size) + scaled = ( + vocab_parallel_logits[:, chunk_start:chunk_end, :].to(work_dtype) + / temperature + ) + logits_max = torch.amax(scaled, dim=-1) + torch.distributed.all_reduce( + logits_max, + op=torch.distributed.ReduceOp.MAX, + group=tp_group, + ) + scaled.sub_(logits_max.unsqueeze(-1)).exp_() + sum_exp = scaled.sum(dim=-1) + torch.distributed.all_reduce( + sum_exp, + op=torch.distributed.ReduceOp.SUM, + group=tp_group, + ) + scaled.div_(sum_exp.unsqueeze(-1)) + scaled.mul_( + grad_output[:, chunk_start:chunk_end].unsqueeze(-1) / temperature + ) + grad_input[:, chunk_start:chunk_end, :].copy_(scaled) + + return grad_input, None, None, None + + class DistributedLogprobWithSampling(torch.autograd.Function): """Custom autograd function for computing log probabilities with top-k/top-p sampling. @@ -995,7 +1363,7 @@ def from_parallel_logits_to_logprobs( target = torch.nn.functional.pad(target, (0, pad_len), value=0) # Shard the targets by context parallelism - cp_rank = torch.distributed.get_rank(cp_group) + cp_rank = 0 if cp_group is None else torch.distributed.get_rank(cp_group) target = _get_tokens_on_this_cp_rank(target, cp_rank, cp_size, seq_dim=1) if need_top_k_or_top_p_filtering(sampling_params): @@ -1251,6 +1619,38 @@ def _get_tokens_on_this_cp_rank( return ids +def cp_load_balanced_segments( + full_seq_len: int, + cp_rank: int, + cp_size: int, +) -> list[tuple[int, int, int]]: + """Describe a CP rank's load-balanced rows as ``(local_start, global_start, length)``. + + With context parallelism, rank ``r`` holds blocks ``r`` and ``2*cp-1-r`` of + a ``2*cp`` partition of the global sequence, in that order (the layout of + :func:`_get_tokens_on_this_cp_rank`). The compact interval metadata is safe + to place in an IPC handle and avoids treating a CP shard as one contiguous + global sequence range. + """ + if cp_size < 1: + raise ValueError(f"cp_size must be >= 1, got {cp_size}") + if not 0 <= cp_rank < cp_size: + raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}") + if cp_size == 1: + return [(0, 0, full_seq_len)] + if full_seq_len % (2 * cp_size) != 0: + raise ValueError( + "Load-balanced context-parallel rows require full_seq_len to be " + f"divisible by 2*cp_size={2 * cp_size}; got {full_seq_len}." + ) + block = full_seq_len // (2 * cp_size) + mirrored = 2 * cp_size - cp_rank - 1 + return [ + (0, cp_rank * block, block), + (block, mirrored * block, block), + ] + + def allgather_cp_sharded_tensor( tensor, cp_group, seq_dim=1 ): # , unpadded_seqlen=None): @@ -1299,6 +1699,54 @@ def backward( # pyrefly: ignore[bad-override] return grad_out, None +class _AllReduceSumWithBackwardSum(torch.autograd.Function): + """SUM all-reduce whose backward also SUM-reduces replicated gradients. + + This variant is for a replicated full-sequence loss under context + parallelism. Every CP rank evaluates the same numerator and the outer loss + wrapper divides by the CP size; summing the backward signal restores the + full gradient on the one rank that owned each sequence value. + """ + + @staticmethod + def forward( # pyrefly: ignore[bad-override] + ctx: Any, + x: torch.Tensor, + group: Optional[torch.distributed.ProcessGroup], + ) -> torch.Tensor: + should_reduce = ( + group is not None + and torch.distributed.is_initialized() + and torch.distributed.get_world_size(group) > 1 + ) + ctx.group = group + ctx.should_reduce = should_reduce + if not should_reduce: + return x + out = x.clone() + torch.distributed.all_reduce( + out, + op=torch.distributed.ReduceOp.SUM, + group=group, + ) + return out + + @staticmethod + def backward( # pyrefly: ignore[bad-override] + ctx: Any, + grad_out: torch.Tensor, + ) -> tuple[torch.Tensor, None]: + if not ctx.should_reduce: + return grad_out, None + grad_input = grad_out.clone() + torch.distributed.all_reduce( + grad_input, + op=torch.distributed.ReduceOp.SUM, + group=ctx.group, + ) + return grad_input, None + + def group_all_reduce_sum_with_grad( x: torch.Tensor, group: Optional[torch.distributed.ProcessGroup], @@ -1313,6 +1761,21 @@ def group_all_reduce_sum_with_grad( return _AllReduceSum.apply(x, group) +def group_all_reduce_sum_with_grad_backward_sum( + x: torch.Tensor, + group: Optional[torch.distributed.ProcessGroup], +) -> torch.Tensor: + """Differentiable SUM whose backward also SUMs over ``group``. + + Use this only when every group rank computes a replica of the same + full-domain numerator and an outer ``1 / world_size`` loss scale must be + cancelled in backward. Tensor-parallel vocabulary reductions should use + :func:`group_all_reduce_sum_with_grad` instead, because their backward must + remain rank-local. + """ + return _AllReduceSumWithBackwardSum.apply(x, group) + + def group_all_reduce_sum( x: torch.Tensor, group: Optional[torch.distributed.ProcessGroup], @@ -1385,8 +1848,9 @@ def backward(ctx, grad_output): ) index = torch.tensor( - [cp_rank, (2 * cp_size - cp_rank - 1)], device="cpu", pin_memory=True - ).cuda(non_blocking=True) + [cp_rank, (2 * cp_size - cp_rank - 1)], + device=grad_output.device, + ) grad_input = grad_output.index_select(seq_dim, index) grad_input = grad_input.view( @@ -1396,6 +1860,88 @@ def backward(ctx, grad_output): return grad_input, None, None # , None +class AllGatherCPContiguousTensor(torch.autograd.Function): + """Autograd-aware all-gather of a *contiguously* CP-sharded tensor. + + Unlike :class:`AllGatherCPTensor` (which un-does PyTorch's load-balanced + ``2*cp`` interleaving), this gathers plain contiguous blocks: CP rank ``r`` + holds ``x[:, r*L:(r+1)*L]`` and the output is the rank-ordered concatenation + ``[B, cp*L, ...]``. Forward concatenates the gathered blocks; backward + all-reduces (SUM) the upstream grad over the CP group and narrows to this + rank's block. So when every CP rank materializes the full sequence and + computes the same scalar, the all-reduce routes the correct full-magnitude + gradient to each rank's own shard (the CP normalization factor cancels), and + every rank participates in the backward collective (no deadlock). + """ + + @staticmethod + def forward( # pyrefly: ignore[bad-override] + ctx: Any, + tensor: torch.Tensor, + cp_group: torch.distributed.ProcessGroup, + seq_dim: int = 1, + ) -> torch.Tensor: + cp_size = torch.distributed.get_world_size(cp_group) + parts = [torch.empty_like(tensor) for _ in range(cp_size)] + torch.distributed.all_gather(parts, tensor.contiguous(), group=cp_group) + ctx.seq_dim = seq_dim + ctx.cp_group = cp_group + ctx.cp_rank = torch.distributed.get_rank(cp_group) + ctx.local_len = tensor.shape[seq_dim] + return torch.cat(parts, dim=seq_dim) + + @staticmethod + def backward( # pyrefly: ignore[bad-override] + ctx: Any, grad_output: torch.Tensor + ) -> tuple[torch.Tensor, None, None]: + torch.distributed.all_reduce(grad_output, group=ctx.cp_group) + grad_input = grad_output.narrow( + ctx.seq_dim, ctx.cp_rank * ctx.local_len, ctx.local_len + ) + return grad_input.contiguous(), None, None + + +def allgather_cp_contiguous_tensor( + tensor: torch.Tensor, + cp_group: Optional[torch.distributed.ProcessGroup], + seq_dim: int = 1, +) -> torch.Tensor: + """All-gather a contiguously CP-sharded tensor to the full sequence. + + No-op (returns the input) when ``cp_group`` is ``None`` or world size <= 1, + so single-process / CPU paths are unchanged. See + :class:`AllGatherCPContiguousTensor` for the load-balanced-vs-contiguous + distinction and the backward semantics. + """ + if cp_group is None or torch.distributed.get_world_size(cp_group) <= 1: + return tensor + return AllGatherCPContiguousTensor.apply(tensor, cp_group, seq_dim) + + +def cp_load_balanced_to_contiguous( + x: torch.Tensor, + *, + cp_group: Optional[torch.distributed.ProcessGroup] = None, + seq_dim: int = 1, +) -> torch.Tensor: + """Re-layout a tensor from load-balanced CP order to this rank's contiguous window. + + PyTorch ``context_parallel`` shards the sequence in a load-balanced + (``2*cp`` interleaved) order. :func:`allgather_cp_sharded_tensor` undoes the + chunking to the full contiguous sequence; this re-slices to this CP rank's + contiguous ``[cp_rank*L, (cp_rank+1)*L)`` window. No-op when CP world <= 1. + Uses the grad-preserving ``DTensor.to_local()`` so the gradient is kept. + """ + if cp_group is None or torch.distributed.get_world_size(cp_group) <= 1: + return x + local = x.to_local() if isinstance(x, DTensor) else x + full = allgather_cp_sharded_tensor(local, cp_group, seq_dim=seq_dim) + cp_size = torch.distributed.get_world_size(cp_group) + cp_rank = torch.distributed.get_rank(cp_group) + local_len = full.shape[seq_dim] // cp_size + return full.narrow(seq_dim, cp_rank * local_len, local_len).contiguous() + + def cp_shift_next( x: torch.Tensor, cp_group: Optional[torch.distributed.ProcessGroup] = None, @@ -1444,6 +1990,72 @@ def vocab_parallel_log_softmax( return torch.log_softmax(scaled, dim=-1) +def chunked_vocab_parallel_logsumexp( + logits: torch.Tensor, + temperature: float, + *, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + chunk_size: int = 256, +) -> torch.Tensor: + """Compute exact full-vocabulary logsumexp without gathering TP logits. + + The result keeps the local sequence layout and drops only the vocabulary + axis. At TP>1 the forward and backward are recomputed in fixed sequence + chunks; at TP1 this uses ordinary ``torch.logsumexp`` in the same chunks so + the unsharded arithmetic remains unchanged. + + Args: + logits: ``[B, S_local, V_local]`` logits, optionally vocab-sharded. + temperature: Positive softmax temperature. + tp_group: Tensor-parallel process group, or ``None`` for TP1. + chunk_size: Positive sequence-axis work bound. + + Returns: + Full-vocabulary log-normalizers with shape ``[B, S_local]``. + """ + if temperature <= 0.0: + raise ValueError(f"temperature must be > 0, got {temperature}") + if chunk_size <= 0: + raise ValueError(f"chunk_size must be > 0, got {chunk_size}") + if isinstance(logits, DTensor): + logits = logits.to_local() + if logits.ndim != 3: + raise ValueError( + "vocab-parallel logsumexp expects [B, S, V] logits, got " + f"{tuple(logits.shape)}" + ) + + tp_size = ( + torch.distributed.get_world_size(tp_group) + if tp_group is not None and torch.distributed.is_initialized() + else 1 + ) + if tp_size > 1: + assert tp_group is not None + return ChunkedVocabParallelLogsumexp.apply( + logits, + float(temperature), + int(chunk_size), + tp_group, + ) + + work_dtype = ( + torch.float32 + if logits.dtype in (torch.float16, torch.bfloat16) + else logits.dtype + ) + chunks = [ + torch.logsumexp( + logits[:, start : start + chunk_size, :].to(work_dtype) / temperature, + dim=-1, + ) + for start in range(0, logits.shape[1], chunk_size) + ] + if not chunks: + return torch.empty((logits.shape[0], 0), dtype=work_dtype, device=logits.device) + return torch.cat(chunks, dim=1) + + def vocab_parallel_full_log_softmax( logits: torch.Tensor, temperature: float, @@ -1522,6 +2134,33 @@ def vocab_parallel_gather_columns( ) +def vocab_parallel_gather_logits( + logits: torch.Tensor, + *, + tp_group: Optional[torch.distributed.ProcessGroup] = None, +) -> torch.Tensor: + """Gather TP vocab-sharded logits to the full vocab ``[..., V]`` (autograd-aware). + + Returns the raw logits over the full vocab so callers that index arbitrary + vocab columns (and compute their own ``logsumexp``) can operate as if the + vocab were unsharded. With ``tp_group`` world > 1 the shards are all-gathered + via Megatron's autograd-aware gather; otherwise the (already-full) local + logits are returned. Unlike :func:`vocab_parallel_full_log_softmax`, no + ``log_softmax`` is applied — use this when the downstream math needs raw + logits rather than log-probs. + """ + if isinstance(logits, DTensor): + logits = logits.to_local() + if tp_group is not None and torch.distributed.get_world_size(tp_group) > 1: + # Local import keeps the optional Megatron dependency boundary intact. + from megatron.core.tensor_parallel import ( + gather_from_tensor_model_parallel_region, + ) + + return gather_from_tensor_model_parallel_region(logits, group=tp_group) + return logits + + def vocab_parallel_argmax( logits: torch.Tensor, *, @@ -1554,6 +2193,8 @@ def get_logprobs_from_vocab_parallel_logits( input_ids: torch.Tensor | DTensor, chunk_size: Optional[int] = None, sampling_params: Optional[TrainingSamplingParams] = None, + *, + seq_index: Optional[torch.Tensor] = None, ): """Computes log probabilities from vocabulary-parallel logits. @@ -1568,22 +2209,38 @@ def get_logprobs_from_vocab_parallel_logits( function will shift them internally. chunk_size (Optional[int]): Sequence dimension chunk size for computing log probabilities. sampling_params (TrainingSamplingParams, optional): Sampling parameters for Top-k/Top-p filtering and temperature scaling. + seq_index: Optional canonical-position permutation for explicitly + CP-sharded DTensor targets. Model-owned CP uses its sharder instead. Returns: torch.Tensor: Log probabilities for the given input IDs, with shape [batch_size, seq_len-1]. The sequence dimension is reduced by 1 due to the target shifting. """ - # This function handles TP only. Fail loudly rather than silently mis-shifting. - if ( + cp_targets = ( isinstance(input_ids, DTensor) and input_ids.device_mesh.mesh_dim_names is not None and "cp" in input_ids.device_mesh.mesh_dim_names - ): + ) + if cp_targets and seq_index is None: raise NotImplementedError( - "get_logprobs_from_vocab_parallel_logits handles TP only; use " - "get_cp_sharded_next_token_logprobs for context parallelism." + "get_logprobs_from_vocab_parallel_logits requires an explicit " + "seq_index for CP DTensor targets; use " + "get_cp_sharded_next_token_logprobs for model-owned context parallelism." ) + if seq_index is not None and not cp_targets: + raise ValueError("seq_index requires CP-sharded DTensor input_ids.") + targets = input_ids.roll(shifts=-1, dims=-1) + if cp_targets: + assert seq_index is not None + target_mesh = input_ids.device_mesh + target_placements = input_ids.placements + sorted_positions = torch.argsort(seq_index) + canonical_targets = input_ids.full_tensor()[:, sorted_positions] + shifted_targets = canonical_targets.roll(shifts=-1, dims=-1)[:, seq_index] + targets = distribute_tensor( + shifted_targets, target_mesh, target_placements + ).to_local() tp_group = vocab_parallel_logits.device_mesh.get_group("tp") tp_rank = tp_group.rank() @@ -1592,7 +2249,7 @@ def get_logprobs_from_vocab_parallel_logits( logprobs = _tp_target_logprobs( vocab_parallel_logits.to_local(), - input_ids.roll(shifts=-1, dims=-1), + targets, vocab_start_index=vocab_interval_per_rank * tp_rank, vocab_end_index=(tp_rank + 1) * vocab_interval_per_rank, tp_group=tp_group, @@ -1601,6 +2258,10 @@ def get_logprobs_from_vocab_parallel_logits( inference_only=not torch.is_grad_enabled(), ) + if cp_targets: + logprobs = DTensor.from_local( + logprobs, target_mesh, target_placements + ).full_tensor()[:, sorted_positions] return logprobs[:, :-1] @@ -1613,6 +2274,8 @@ def get_next_token_logprobs_from_logits( sampling_params: Optional[TrainingSamplingParams] = None, chunk_size: Optional[int] = None, cp_sharder: Optional["ContextParallelSharder"] = None, + *, + seq_index: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Compute token log-probabilities from logits, handling parallel and non-parallel cases. @@ -1635,6 +2298,8 @@ def get_next_token_logprobs_from_logits( forward's model batch; set only when cp_size > 1. When set, ``next_token_logits`` is this rank's CP-local shard and the sharder owns the sequence layout. + seq_index: Optional canonical-position permutation for explicitly + CP-sharded DTensor targets without a model-owned sharder. Returns: Token log-probabilities of shape [batch_size, seq_len - 1] @@ -1680,6 +2345,7 @@ def get_next_token_logprobs_from_logits( next_token_logits, input_ids, sampling_params=sampling_params, + seq_index=seq_index, ) else: @@ -1774,6 +2440,324 @@ def distributed_vocab_topk( return topk_vals, topk_global_indices +def _canonical_topk_by_value_then_lowest_id( + values: torch.Tensor, + global_ids: torch.Tensor, + k: int, + *, + id_ceiling: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Select top-k by value with the lowest global token ID as tie-breaker. + + ``torch.topk`` does not specify which equal-valued entries it returns, and + its candidate order changes when the vocabulary is partitioned across TP. + Build an exact integer lexicographic key from the fp32 IEEE-754 ordering and + the global token ID instead. This preserves every non-tied comparison while + making exact cutoff ties topology-independent. ``global_ids`` may be a + one-dimensional vocabulary vector or have the same shape as ``values``. + """ + if values.is_complex() or not values.is_floating_point(): + raise TypeError("canonical top-k values must be floating point") + if global_ids.dtype == torch.bool or global_ids.is_floating_point(): + raise TypeError("canonical top-k token ids must have an integer dtype") + if k <= 0 or k > values.shape[-1]: + raise ValueError( + f"canonical top-k requires 0 < k <= {values.shape[-1]}, got {k}" + ) + if id_ceiling < 0 or id_ceiling >= (1 << 31): + raise ValueError( + f"canonical top-k requires id_ceiling in [0, 2**31), got {id_ceiling}" + ) + + scores = values.float() + ids = global_ids.to(device=values.device, dtype=torch.int64) + try: + scores, ids = torch.broadcast_tensors(scores, ids) + except RuntimeError as exc: + raise ValueError( + "canonical top-k values and token ids are not broadcastable: " + f"{tuple(values.shape)} vs {tuple(global_ids.shape)}" + ) from exc + if bool(torch.isnan(scores).any().item()): + raise ValueError("canonical top-k does not accept NaN logits") + if bool(((ids < 0) | (ids > int(id_ceiling))).any().item()): + raise ValueError(f"canonical top-k token ids must be in [0, {id_ceiling}]") + + # Normalize negative zero so numerically equal +/-0 values use only the ID + # tie-break. BF16/FP16 values are represented exactly after the fp32 cast. + scores = torch.where(scores == 0, torch.zeros_like(scores), scores).contiguous() + raw_bits = scores.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + sign_set = (raw_bits & 0x80000000) != 0 + ordered_score = torch.where( + sign_set, + (~raw_bits) & 0xFFFFFFFF, + raw_bits ^ 0x80000000, + ) + id_bits = max(1, int(id_ceiling).bit_length()) + secondary = int(id_ceiling) - ids + keys = (ordered_score << id_bits) + secondary + positions = torch.topk(keys, k, dim=-1, largest=True, sorted=True).indices + return ( + torch.gather(scores, dim=-1, index=positions), + torch.gather(ids, dim=-1, index=positions), + ) + + +@torch.no_grad() +def distributed_vocab_topk_logz_force( + vocab_parallel_logits: torch.Tensor, + force_token_ids: torch.Tensor, + k: int, + tp_group: Optional[torch.distributed.ProcessGroup], + *, + vocab_start_index: int, + vocab_end_index: int, + real_vocab_size: int, + temperature: float, + noise_filter_k: int = 0, + chunk_size: Optional[int] = None, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Teacher sparse transport primitive: natural top-k plus forced sidecars. + + The vocabulary remains TP-sharded. Each rank contributes at most + ``max(k, noise_filter_k)`` local candidates and row-wise MAX/SUM scalars; + no full vocabulary tensor is gathered. ``force_token_ids`` has shape + ``[B, S, F]`` and uses ``-1`` for unused sidecar slots. Every valid forced + label is resolved across TP without modifying the natural top-k. + + Returns, in order: + + * raw (unscaled) natural top-k logits sorted by token id; + * their global int32 token ids; + * exact ``logsumexp(logits / temperature)``; + * the natural score-rank-k token id, used to reproduce exact cutoff-tie + eviction after the natural row has been sorted by id; + * raw forced-label logits; + * forced global int32 ids (including ``-1`` padding); and + * per-forced-label membership in natural top-``noise_filter_k`` (or natural + top-k when the noise filter is disabled). + + The consumer constructs each loss request independently as natural K when + its required label is already present, otherwise natural K-1 plus that + request's forced label. Keeping the sidecar separate is necessary when two + aligned chunks request different labels from the same teacher row. + """ + if vocab_parallel_logits.dim() != 3: + raise ValueError( + "vocab_parallel_logits must have shape [B, S, V_local], got " + f"{tuple(vocab_parallel_logits.shape)}" + ) + if ( + force_token_ids.dim() != 3 + or force_token_ids.shape[:2] != vocab_parallel_logits.shape[:2] + ): + raise ValueError( + "force_token_ids must have shape [B, S, F] matching logits, got " + f"{tuple(force_token_ids.shape)} vs " + f"{tuple(vocab_parallel_logits.shape[:2])}" + ) + force_width = int(force_token_ids.shape[-1]) + if force_width <= 0: + raise ValueError("force_token_ids sidecar width F must be positive") + if force_token_ids.dtype == torch.bool or force_token_ids.is_floating_point(): + raise TypeError("force_token_ids must have an integer dtype") + if not math.isfinite(float(temperature)) or temperature <= 0: + raise ValueError(f"temperature must be finite and > 0, got {temperature}") + if k <= 0: + raise ValueError(f"k must be > 0, got {k}") + if noise_filter_k < 0: + raise ValueError(f"noise_filter_k must be >= 0, got {noise_filter_k}") + k_compute = max(int(k), int(noise_filter_k)) + if k_compute > int(real_vocab_size): + raise ValueError( + f"max(k, noise_filter_k)={k_compute} exceeds real_vocab_size=" + f"{real_vocab_size}" + ) + invalid_force = (force_token_ids < -1) | (force_token_ids >= int(real_vocab_size)) + if bool(invalid_force.any().item()): + bad_id = int(force_token_ids[invalid_force][0].item()) + raise ValueError( + "force_token_ids entries must be -1 or real-vocabulary ids in " + f"[0, {real_vocab_size}); got {bad_id}." + ) + for left in range(force_width): + for right in range(left + 1, force_width): + duplicate = (force_token_ids[..., left] >= 0) & ( + force_token_ids[..., left] == force_token_ids[..., right] + ) + if bool(duplicate.any().item()): + raise ValueError( + "force_token_ids contains duplicate valid labels in one row" + ) + + tp_world = ( + torch.distributed.get_world_size(tp_group) + if tp_group is not None and torch.distributed.is_initialized() + else 1 + ) + B, S, V_local = vocab_parallel_logits.shape + local_real_count = max( + 0, + min(int(vocab_end_index), int(real_vocab_size)) - int(vocab_start_index), + ) + local_real_count = min(local_real_count, V_local) + if chunk_size is None or chunk_size <= 0: + chunk_size = S + + out_values: list[torch.Tensor] = [] + out_indices: list[torch.Tensor] = [] + out_log_z: list[torch.Tensor] = [] + out_natural_tail_indices: list[torch.Tensor] = [] + out_forced_values: list[torch.Tensor] = [] + out_forced_indices: list[torch.Tensor] = [] + out_forced_in_topk: list[torch.Tensor] = [] + sentinel = int(real_vocab_size) + membership_k = int(noise_filter_k) if noise_filter_k > 0 else int(k) + + for s0 in range(0, S, chunk_size): + s1 = min(S, s0 + chunk_size) + raw = vocab_parallel_logits[:, s0:s1, :].float() + rows_shape = raw.shape[:2] + + # Exact real-vocabulary logZ. Padded LM-head columns never enter the + # local MAX or SUM; a rank with no real columns contributes (-inf, 0). + if local_real_count > 0: + scaled_real = raw[..., :local_real_count] / float(temperature) + local_max = scaled_real.max(dim=-1).values + else: + scaled_real = raw[..., :0] + local_max = torch.full( + rows_shape, -torch.inf, dtype=torch.float32, device=raw.device + ) + if tp_world > 1: + torch.distributed.all_reduce( + local_max, op=torch.distributed.ReduceOp.MAX, group=tp_group + ) + if local_real_count > 0: + local_sum = torch.exp(scaled_real - local_max.unsqueeze(-1)).sum(dim=-1) + else: + local_sum = torch.zeros(rows_shape, dtype=torch.float32, device=raw.device) + if tp_world > 1: + torch.distributed.all_reduce( + local_sum, op=torch.distributed.ReduceOp.SUM, group=tp_group + ) + log_z = local_max + local_sum.log() + + # Fixed-width candidate buffers avoid selecting padded/sentinel ids on + # a TP rank that owns fewer than k_compute real tokens. + local_vals = torch.full( + (*rows_shape, k_compute), + -torch.inf, + dtype=torch.float32, + device=raw.device, + ) + local_ids = torch.full( + (*rows_shape, k_compute), + sentinel, + dtype=torch.long, + device=raw.device, + ) + local_k = min(k_compute, local_real_count) + if local_k > 0: + local_global_ids = torch.arange( + int(vocab_start_index), + int(vocab_start_index) + local_real_count, + dtype=torch.long, + device=raw.device, + ) + vals, ids = _canonical_topk_by_value_then_lowest_id( + raw[..., :local_real_count], + local_global_ids, + local_k, + id_ceiling=int(real_vocab_size), + ) + local_vals[..., :local_k] = vals + local_ids[..., :local_k] = ids + + if tp_world > 1: + gathered_vals = [torch.empty_like(local_vals) for _ in range(tp_world)] + gathered_ids = [torch.empty_like(local_ids) for _ in range(tp_world)] + torch.distributed.all_gather(gathered_vals, local_vals, group=tp_group) + torch.distributed.all_gather(gathered_ids, local_ids, group=tp_group) + candidate_vals = torch.cat(gathered_vals, dim=-1) + candidate_ids = torch.cat(gathered_ids, dim=-1) + else: + candidate_vals = local_vals + candidate_ids = local_ids + natural_vals, natural_ids = _canonical_topk_by_value_then_lowest_id( + candidate_vals, + candidate_ids, + k_compute, + id_ceiling=int(real_vocab_size), + ) + if bool((natural_ids >= int(real_vocab_size)).any().item()): + raise RuntimeError( + "Distributed top-k selected a padded/sentinel vocabulary id; " + "check the real-vocabulary candidate accounting." + ) + + force = force_token_ids[:, s0:s1, :].long() + force_valid = (force >= 0) & (force < int(real_vocab_size)) + forced_in_topk = ( + natural_ids[..., :membership_k].unsqueeze(-2) == force.unsqueeze(-1) + ).any(dim=-1) & force_valid + + owned = ( + force_valid + & (force >= int(vocab_start_index)) + & (force < int(vocab_start_index) + local_real_count) + ) + local_force_id = (force - int(vocab_start_index)).clamp( + min=0, max=max(V_local - 1, 0) + ) + forced_values = torch.gather(raw, -1, local_force_id) + forced_values = torch.where( + owned, forced_values, torch.zeros_like(forced_values) + ) + if tp_world > 1: + torch.distributed.all_reduce( + forced_values, op=torch.distributed.ReduceOp.SUM, group=tp_group + ) + forced_values = torch.where( + force_valid, forced_values, torch.zeros_like(forced_values) + ) + + top_vals = natural_vals[..., :k] + top_ids = natural_ids[..., :k] + natural_tail_ids = top_ids[..., -1:].clone() + + # The sparse lookup uses searchsorted, so canonicalize by token id + # after the score-ranked natural selection. The pre-sort tail id is + # transported separately so cutoff-logit ties keep score-ranked eviction. + sorted_ids, order = torch.sort(top_ids, dim=-1) + sorted_vals = torch.gather(top_vals, dim=-1, index=order) + out_values.append(sorted_vals) + out_indices.append(sorted_ids.to(torch.int32)) + out_log_z.append(log_z) + out_natural_tail_indices.append(natural_tail_ids.to(torch.int32)) + out_forced_values.append(forced_values) + out_forced_indices.append(force.to(torch.int32)) + out_forced_in_topk.append(forced_in_topk) + + return ( + torch.cat(out_values, dim=1), + torch.cat(out_indices, dim=1), + torch.cat(out_log_z, dim=1), + torch.cat(out_natural_tail_indices, dim=1), + torch.cat(out_forced_values, dim=1), + torch.cat(out_forced_indices, dim=1), + torch.cat(out_forced_in_topk, dim=1), + ) + + def gather_logits_at_global_indices( vocab_parallel_logits: torch.Tensor, global_indices: torch.Tensor, diff --git a/nemo_rl/models/automodel/config.py b/nemo_rl/models/automodel/config.py index ba959aa3345..fe522ce33f5 100644 --- a/nemo_rl/models/automodel/config.py +++ b/nemo_rl/models/automodel/config.py @@ -72,6 +72,11 @@ class RuntimeConfig(NamedTuple): # Reward model flag is_reward_model: bool + # dtype the checkpoint is materialized in (automodel_cfg.load_precision). + # Trails the required fields with a default so existing keyword + # constructions stay valid. + model_load_dtype: Optional[torch.dtype] = None + class ModelAndOptimizerState(NamedTuple): """Container for model and optimizer state. diff --git a/nemo_rl/models/automodel/data.py b/nemo_rl/models/automodel/data.py index 3aea08fa75d..ec799fb3b0f 100644 --- a/nemo_rl/models/automodel/data.py +++ b/nemo_rl/models/automodel/data.py @@ -24,12 +24,14 @@ from torch import nn from transformers import AutoTokenizer +from nemo_rl.algorithms.loss.csd_packing import resolve_lockstep_bin_geometry from nemo_rl.algorithms.loss.interfaces import LossFunction, LossType from nemo_rl.data.multimodal_utils import PACKED_MULTIMODAL_FIELDS from nemo_rl.data_plane.codec import pad_field from nemo_rl.data_plane.schema import MICROBATCH_PADDED_FIELDS from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.huggingface.common import ( + FlashAttentionKwargs, get_flash_attention_kwargs, pack_sequences, ) @@ -103,6 +105,101 @@ def filter_multimodal_kwargs_for_model( } +@dataclass +class PaddedFlashAttentionKwargs(FlashAttentionKwargs): + """Flash-attention metadata with distinct logical and physical offsets.""" + + cu_seqlens_q_padded: torch.Tensor + cu_seqlens_k_padded: torch.Tensor + + @property + def cu_seq_lens_q(self) -> torch.Tensor: + """Physical segment boundaries consumed by HF flash-attention.""" + return self.cu_seqlens_q_padded + + @property + def cu_seq_lens_k(self) -> torch.Tensor: + """Physical segment boundaries consumed by HF flash-attention.""" + return self.cu_seqlens_k_padded + + @property + def max_length_q(self) -> int: + return self.max_seqlen_q + + @property + def max_length_k(self) -> int: + return self.max_seqlen_k + + +def _materialize_planned_packed_inputs( + *, + input_ids: torch.Tensor, + input_lengths: torch.Tensor, + raw_cu_seqlens: tuple[int, ...], + padded_cu_seqlens: tuple[int, ...], + physical_tokens: int, + padding_value: int, +) -> tuple[torch.Tensor, torch.Tensor, PaddedFlashAttentionKwargs]: + """Materialize exactly one controller-planned packed row.""" + if len(raw_cu_seqlens) != input_ids.shape[0] + 1: + raise ValueError( + "Planned raw cumulative boundaries do not match the logical batch " + f"size {input_ids.shape[0]}." + ) + if len(padded_cu_seqlens) != len(raw_cu_seqlens): + raise ValueError("Planned raw and padded boundaries have different shapes.") + if padded_cu_seqlens[-1] != physical_tokens: + raise ValueError( + f"Planned padded boundaries end at {padded_cu_seqlens[-1]}, not " + f"physical size {physical_tokens}." + ) + + token_segments: list[torch.Tensor] = [] + position_segments: list[torch.Tensor] = [] + for row in range(input_ids.shape[0]): + raw_length = int(input_lengths[row].item()) + padded_length = padded_cu_seqlens[row + 1] - padded_cu_seqlens[row] + if raw_length > input_ids.shape[1] or padded_length < raw_length: + raise ValueError( + f"Cannot materialize planned row {row}: raw_length={raw_length}, " + f"padded_length={padded_length}, source_width={input_ids.shape[1]}." + ) + token_segment = input_ids[row, :raw_length] + position_segment = torch.arange( + raw_length, dtype=torch.long, device=input_ids.device + ) + if padded_length > raw_length: + padding = padded_length - raw_length + token_segment = torch.nn.functional.pad( + token_segment, (0, padding), value=padding_value + ) + position_segment = torch.nn.functional.pad(position_segment, (0, padding)) + token_segments.append(token_segment) + position_segments.append(position_segment) + + packed_input_ids = torch.cat(token_segments).unsqueeze(0) + packed_position_ids = torch.cat(position_segments).unsqueeze(0) + if packed_input_ids.shape[1] != physical_tokens: + raise ValueError( + "Materialized lockstep physical size drifted: got " + f"{packed_input_ids.shape[1]}, planned {physical_tokens}." + ) + raw_cu = torch.tensor(raw_cu_seqlens, dtype=torch.int32, device=input_ids.device) + padded_cu = torch.tensor( + padded_cu_seqlens, dtype=torch.int32, device=input_ids.device + ) + physical_lengths = padded_cu[1:] - padded_cu[:-1] + flash_kwargs = PaddedFlashAttentionKwargs( + cu_seqlens_q=raw_cu, + cu_seqlens_k=raw_cu.clone(), + max_seqlen_q=int(physical_lengths.max().item()), + max_seqlen_k=int(physical_lengths.max().item()), + cu_seqlens_q_padded=padded_cu, + cu_seqlens_k_padded=padded_cu.clone(), + ) + return packed_input_ids, packed_position_ids, flash_kwargs + + @dataclass class ProcessedInputs: """Processed microbatch inputs ready for model forward pass. @@ -285,23 +382,59 @@ def process_microbatch( input_ids = mb.get("input_ids").cuda() if enable_seq_packing: - input_ids, position_ids, _ = pack_sequences( - input_ids=input_ids, - input_lengths=mb["input_lengths"], - packed_sequence_size=[ - len(mb["input_lengths"]) - ], # flash attention 2 expects flattened input - padding_value=tokenizer.eos_token_id, - return_attention_mask=False, - min_seq_len=cfg["sequence_packing"][ - "train_mb_tokens" - ], # TODO: this is a WAR for sequence packing, we should fix this. Without this, backward will fail when TP is enabled. + planned_geometry = resolve_lockstep_bin_geometry( + mb, input_lengths=mb["input_lengths"] ) + if planned_geometry is not None: + cp_size = int(cfg["automodel_cfg"]["context_parallel_size"]) + if cp_size != 1: + raise NotImplementedError( + "Plan-driven Automodel packing currently supports only " + "context_parallel_size=1." + ) + tp_size = int(cfg["automodel_cfg"]["tensor_parallel_size"]) + train_mb_tokens = int(cfg["sequence_packing"]["train_mb_tokens"]) + if ( + getattr(mb, "lockstep_side_id", None) == "student" + and tp_size > 1 + and planned_geometry.physical_tokens != train_mb_tokens + ): + raise ValueError( + "Plan-driven Automodel training with tensor_parallel_size>1 " + "requires every student bin to preserve the fixed " + "train_mb_tokens tail used by generic sequence packing: " + f"planned={planned_geometry.physical_tokens}, " + f"train_mb_tokens={train_mb_tokens}. Configure the student " + "SidePackingSpec physical_size_fn to return train_mb_tokens." + ) + input_ids, position_ids, flash_attn_kwargs = ( + _materialize_planned_packed_inputs( + input_ids=input_ids, + input_lengths=mb["input_lengths"], + raw_cu_seqlens=planned_geometry.raw_cu_seqlens, + padded_cu_seqlens=planned_geometry.padded_cu_seqlens, + physical_tokens=planned_geometry.physical_tokens, + padding_value=tokenizer.eos_token_id, + ) + ) + else: + input_ids, position_ids, _ = pack_sequences( + input_ids=input_ids, + input_lengths=mb["input_lengths"], + packed_sequence_size=[ + len(mb["input_lengths"]) + ], # flash attention 2 expects flattened input + padding_value=tokenizer.eos_token_id, + return_attention_mask=False, + min_seq_len=cfg["sequence_packing"][ + "train_mb_tokens" + ], # TODO: this is a WAR for sequence packing, we should fix this. Without this, backward will fail when TP is enabled. + ) + flash_attn_kwargs = get_flash_attention_kwargs( + input_lengths=mb["input_lengths"], + ) seq_len = input_ids.shape[1] attention_mask = None - flash_attn_kwargs = get_flash_attention_kwargs( - input_lengths=mb["input_lengths"], - ) else: batch_size, seq_len = input_ids.shape @@ -367,6 +500,8 @@ def process_global_batch( - batch: The extracted batch - global_valid_seqs: Number of valid sequences across all ranks - global_valid_toks: Number of valid tokens across all ranks + - global_valid_chunks_by_idx: Number of valid cross-tokenizer alignment + chunks for each teacher across all data-parallel ranks """ batch = data.get_batch(batch_idx=batch_idx, batch_size=batch_size) @@ -381,10 +516,41 @@ def process_global_batch( local_valid_toks = torch.sum( batch["token_mask"][:, 1:] * batch["sample_mask"].unsqueeze(-1) ) + kd_token_mask = batch.get("kd_token_mask", batch.get("token_mask")) + if kd_token_mask is None: + local_valid_kd_toks = local_valid_toks + else: + local_valid_kd_toks = torch.sum( + kd_token_mask[:, 1:] * batch["sample_mask"].unsqueeze(-1) + ) - to_reduce = torch.tensor([local_valid_seqs, local_valid_toks]).cuda() + local_valid_chunks_by_idx: dict[int, torch.Tensor] = {} + for key, pair_valid in batch.items(): + if not key.startswith("alignment_") or not key.endswith("_pair_valid"): + continue + teacher_idx = key.removeprefix("alignment_").removesuffix("_pair_valid") + if not teacher_idx.isdigit(): + continue + sample_mask = batch["sample_mask"].to(pair_valid.dtype).unsqueeze(-1) + local_valid_chunks_by_idx[int(teacher_idx)] = torch.sum( + pair_valid * sample_mask + ) + + chunk_teacher_indices = sorted(local_valid_chunks_by_idx) + to_reduce = torch.stack( + [ + local_valid_seqs, + local_valid_toks, + local_valid_kd_toks, + *(local_valid_chunks_by_idx[i] for i in chunk_teacher_indices), + ] + ).cuda() torch.distributed.all_reduce(to_reduce, group=dp_group) - global_valid_seqs, global_valid_toks = to_reduce[0], to_reduce[1] + global_valid_seqs, global_valid_toks, global_valid_kd_toks = to_reduce[:3] + global_valid_chunks_by_idx = { + teacher_idx: to_reduce[offset + 3] + for offset, teacher_idx in enumerate(chunk_teacher_indices) + } if hasattr(loss_fn, "loss_type") and loss_fn.loss_type == LossType.TOKEN_LEVEL: assert "token_mask" in batch, ( @@ -395,6 +561,8 @@ def process_global_batch( "batch": batch, "global_valid_seqs": global_valid_seqs, "global_valid_toks": global_valid_toks, + "global_valid_kd_toks": global_valid_kd_toks, + "global_valid_chunks_by_idx": global_valid_chunks_by_idx, } diff --git a/nemo_rl/models/automodel/nemotron_h_cp.py b/nemo_rl/models/automodel/nemotron_h_cp.py new file mode 100644 index 00000000000..f03f9569d2c --- /dev/null +++ b/nemo_rl/models/automodel/nemotron_h_cp.py @@ -0,0 +1,272 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Context parallelism for the HuggingFace (``force_hf``) NemotronH implementation. + +AUTOMODEL-WORKAROUND(nemotron-h-cp): Automodel's ``NemotronHParallelizationStrategy`` +selects on the class name ``NemotronHForCausalLM``, which both implementations share, +but the HF model needs additional CP setup. The attention branch reads +``layer.mixer.attn_module`` and the Mamba branch sets ``mixer.cp``. These +interfaces belong to the custom implementation. On the HF model the attention +attribute is absent, and its Mamba ``forward`` does not consume ``mixer.cp``. +Without a CP-aware forward, each Mamba layer would process only its local +sequence shard without exchanging recurrent state. + +A further gap sits in the checkpoint's own remote code and fires first: NeMo-RL and +Automodel both request ``attn_implementation="sdpa"`` once ``context_parallel_size +> 1`` (ring attention needs torch SDPA, not flash/TE), but +``NemotronHPreTrainedModel`` never sets ``_supports_sdpa``, so transformers >= 5 +rejects the request outright: + + ValueError: NemotronHForCausalLM does not support an attention implementation + through torch.nn.functional.scaled_dot_product_attention yet. + +The remote code does ship an SDPA path (``NEMOTRONH_ATTENTION_CLASSES["sdpa"] -> +NemotronHSdpaAttention``, and even the "eager" class calls +``F.scaled_dot_product_attention``); only the capability flag is missing. + +This module provides the required adaptations for the HF model, without touching the 3rdparty +submodule or the checkpoint: + +* Attention: the HF mixer is torch ``scaled_dot_product_attention`` (NeMo-RL + forces ``attn_implementation="sdpa"`` whenever ``context_parallel_size > 1``), + which torch's ``context_parallel()`` context manager already ring-patches. The + TE ``set_context_parallel_group`` hook is therefore not needed -- an + ``attn_module = None`` marker makes upstream's ``isinstance`` check skip it. +* Mamba: ``NemotronHMamba2Mixer.forward`` is rebound per instance to a CP-aware + variant that routes the fused ``mamba_split_conv1d_scan_combined`` call through + :class:`MambaContextParallel` (all-to-all to a hidden-parallel layout so each + rank sees the full sequence for a slice of the heads), mirroring the custom + ``nemotron_v3`` mixer's CP path. + +Call :func:`apply_nemotron_h_cp_fixes` once per worker process before the model +is built; it is idempotent and a no-op for the custom implementation. +""" + +from __future__ import annotations + +import sys +import types +from typing import Any, Optional + +import torch + + +def _module_of(obj: Any) -> Any: + """Return the (dynamically loaded) python module a HF remote-code object came from.""" + return sys.modules[type(obj).__module__] + + +def _cp_aware_mamba_forward( + self: Any, + hidden_states: torch.Tensor, + cache_params: Optional[Any] = None, + cache_position: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """CP-aware replacement for ``NemotronHMamba2Mixer.forward``. + + Delegates to the stock forward unless a :class:`MambaContextParallel` helper is + attached and active. The CP path mirrors Automodel's custom ``nemotron_v3`` + mixer: redistribute the ``in_proj`` output from sequence-sharded to + hidden-sharded, run the fused conv+SSM kernel over the *full* sequence with + per-rank head/group slices of ``conv1d`` / ``dt_bias`` / ``A_log`` / ``D``, + redistribute back, then apply the gated RMSNorm and ``out_proj`` (both of + which the single-GPU path fuses into the kernel and which must run on the + full hidden dimension). + """ + cp = getattr(self, "cp", None) + if cp is None or cp.cp_size == 1: + return self._nrl_orig_forward( + hidden_states, cache_params, cache_position, attention_mask + ) + if cache_params is not None: + raise NotImplementedError( + "NemotronH Mamba context parallelism supports the training forward only; " + "cached/step decoding under CP is not implemented." + ) + + # Nemotron-H supplies this optional runtime dependency; other models should + # not have to install it merely to import NeMo RL. + from mamba_ssm.ops.triton.ssd_combined import ( # pyrefly: ignore[import-error] + mamba_split_conv1d_scan_combined, + ) + + hf = _module_of(self) + hidden_states = hf.apply_mask_to_padding_states(hidden_states, attention_mask) + projected_states = self.in_proj(hidden_states) + + dt_limit_kwargs = ( + {} + if self.time_step_limit == (0.0, float("inf")) + else {"dt_limit": self.time_step_limit} + ) + + projected_states = cp.pre_conv_ssm(projected_states) + A = -torch.exp(cp.get_A_log().float()) + out = mamba_split_conv1d_scan_combined( + projected_states, + cp.get_conv1d_weight(), + cp.get_conv1d_bias(), + cp.get_dt_bias(), + A, + D=cp.get_D(), + chunk_size=self.chunk_size, + seq_idx=None, + activation=self.activation, + # rmsnorm/outproj are deliberately NOT fused here: they operate on the + # full hidden dimension, which only exists after post_conv_ssm. With + # rmsnorm_weight=None the kernel still applies the silu gate inside the + # scan, so the composition below stays RMSNorm(x_ssm * silu(z)) -- the + # same expression the fused single-GPU call computes. + rmsnorm_weight=None, + outproj_weight=None, + headdim=self.head_dim, + ngroups=cp.n_groups_local, + norm_before_gate=False, + return_final_states=False, + **dt_limit_kwargs, + ) + if out.ndim == 4: + out = out.reshape(out.shape[0], out.shape[1], -1) + out = cp.post_conv_ssm(out) + out = self.norm(out, gate=None) + return self.out_proj(out) + + +def _bind_cp_aware_mamba_forward(mixer: torch.nn.Module) -> None: + """Rebind ``mixer.forward`` to the CP-aware variant (idempotent).""" + if hasattr(mixer, "_nrl_orig_forward"): + return + object.__setattr__(mixer, "_nrl_orig_forward", mixer.forward) + mixer.forward = types.MethodType(_cp_aware_mamba_forward, mixer) + if not hasattr(mixer, "cp"): + object.__setattr__(mixer, "cp", None) + + +def prepare_hf_nemotron_h_for_cp(model: torch.nn.Module) -> int: + """Make an HF remote-code NemotronH survive and honor Automodel's CP branch. + + Returns the number of Mamba mixers rebound (0 for the custom implementation, + which Automodel's MoE parallelizer handles natively). + """ + backbone = getattr(model, "backbone", None) + if backbone is None: + return 0 + + rebound = 0 + for layer in backbone.layers: + block_type = getattr(layer, "block_type", None) + mixer = getattr(layer, "mixer", None) + if mixer is None: + continue + if block_type == "attention" and not hasattr(mixer, "attn_module"): + # Upstream does `isinstance(layer.mixer.attn_module, DotProductAttention)`; + # the HF mixer has no such attribute. None makes that check False, so the + # TE-only hook is skipped -- correct here, because torch's + # context_parallel() already ring-patches the SDPA this mixer calls. + mixer.attn_module = None + elif block_type == "mamba": + _bind_cp_aware_mamba_forward(mixer) + rebound += 1 + return rebound + + +def _declare_sdpa_support(cls: Any) -> bool: + """Set ``_supports_sdpa`` on a dynamically loaded NemotronH class. + + Returns True if the flag was flipped. Only applied when the loaded module + actually maps an SDPA attention class, so this never claims a capability the + checkpoint's code does not implement. + """ + module = sys.modules.get(getattr(cls, "__module__", ""), None) + attention_classes = getattr(module, "NEMOTRONH_ATTENTION_CLASSES", None) + if not attention_classes or "sdpa" not in attention_classes: + return False + flipped = False + for klass in cls.__mro__: + if klass.__name__ == "NemotronHPreTrainedModel" and not getattr( + klass, "_supports_sdpa", False + ): + klass._supports_sdpa = True + flipped = True + return flipped + + +def _patch_dynamic_module_sdpa_flag() -> None: + """Declare SDPA support on NemotronH classes as transformers loads them. + + ``get_class_in_module`` is where ``trust_remote_code`` classes materialize, and + it is reached from every entry point (``AutoModel``, Automodel's + ``_from_pretrained_parent_class``). Patching it means the flag is set before + ``_sdpa_can_dispatch`` runs, without NeMo-RL having to resolve the remote class + itself. + """ + import transformers.dynamic_module_utils as dynamic_module_utils + + if getattr(dynamic_module_utils, "_nrl_nemotron_h_sdpa_patched", False): + return + + original = dynamic_module_utils.get_class_in_module + + def _get_class_in_module(class_name, *args, **kwargs): + cls = original(class_name, *args, **kwargs) + if isinstance(cls, type) and class_name.startswith("NemotronH"): + if _declare_sdpa_support(cls): + print( + "[nemo-rl] NemotronH (HF): declared _supports_sdpa=True " + "(remote code implements SDPA but omits the flag; " + "context parallelism requires attn_implementation='sdpa')", + flush=True, + ) + return cls + + dynamic_module_utils.get_class_in_module = _get_class_in_module + dynamic_module_utils._nrl_nemotron_h_sdpa_patched = True + + +def apply_nemotron_h_cp_fixes() -> None: + """Install the CP-corrected NemotronH parallelization strategy (idempotent). + + Wraps Automodel's strategy rather than replacing it: the pre-pass runs on the + unparallelized model, then the upstream body (TP plan, ``MambaContextParallel`` + attachment, activation checkpointing, FSDP) runs unchanged, so this stays + correct if upstream's plan evolves. + """ + from nemo_automodel.components.distributed import parallelizer as _p + + _patch_dynamic_module_sdpa_flag() + + if getattr(_p, "_nrl_nemotron_h_cp_applied", False): + return + + upstream_cls = type(_p.PARALLELIZATION_STRATEGIES["NemotronHForCausalLM"]) + + class _NemotronHCPFixStrategy(upstream_cls): # type: ignore[misc,valid-type] + """NemotronH strategy that also supports the HF remote-code model under CP.""" + + def parallelize(self, model, device_mesh, **kwargs): # type: ignore[override] + mesh_names = device_mesh.mesh_dim_names or () + cp_mesh = device_mesh["cp"] if "cp" in mesh_names else None + if cp_mesh is not None and cp_mesh.size() > 1: + rebound = prepare_hf_nemotron_h_for_cp(model) + if rebound: + print( + f"[nemo-rl] NemotronH (HF): enabled Mamba context parallelism " + f"on {rebound} mixers (cp={cp_mesh.size()})", + flush=True, + ) + return super().parallelize(model, device_mesh, **kwargs) + + _p.PARALLELIZATION_STRATEGIES["NemotronHForCausalLM"] = _NemotronHCPFixStrategy() + _p._nrl_nemotron_h_cp_applied = True diff --git a/nemo_rl/models/automodel/setup.py b/nemo_rl/models/automodel/setup.py index ed1711b2f5c..66e3d7d4813 100644 --- a/nemo_rl/models/automodel/setup.py +++ b/nemo_rl/models/automodel/setup.py @@ -18,6 +18,7 @@ import inspect import json import os +from contextlib import contextmanager from functools import partial from typing import Any, Optional, Union @@ -69,6 +70,7 @@ ModelAndOptimizerState, RuntimeConfig, ) +from nemo_rl.models.automodel.nemotron_h_cp import apply_nemotron_h_cp_fixes from nemo_rl.models.automodel.router_replay import ( enable_routing_replay_in_automodel_kwargs, require_automodel_moe_model, @@ -106,6 +108,43 @@ def _has_optimizer_fp32_master( ) +def _precision_to_dtype(precision: str, field_name: str) -> torch.dtype: + if precision not in STRING_TO_DTYPE: + raise ValueError(f"Unknown {field_name}: {precision}") + return STRING_TO_DTYPE[precision] + + +@contextmanager +def _force_shard_before_load(enabled: bool): + """Make Automodel shard the model before loading its weights. + + Automodel decides load-vs-shard order from the parallelism sizes; with + tp_size == 1 it loads the unwrapped model first, which materializes the + whole checkpoint on every rank. Forcing the post-shard path caps peak + memory at the per-rank shard, which is what makes 30B-class teachers fit. + """ + if not enabled: + yield + return + + import nemo_automodel._transformers.infrastructure as automodel_infrastructure + + original_should_load_before_shard = ( + automodel_infrastructure._should_load_before_shard + ) + + def _always_load_after_shard(**_kwargs: Any) -> bool: + return False + + automodel_infrastructure._should_load_before_shard = _always_load_after_shard + try: + yield + finally: + automodel_infrastructure._should_load_before_shard = ( + original_should_load_before_shard + ) + + def _maybe_set_force_hf(automodel_kwargs: dict, model_config) -> None: """Validate and maybe auto-set force_hf based on adapter compatibility. @@ -280,6 +319,7 @@ def validate_and_prepare_config( config: PolicyConfig, processor: Optional[AutoProcessor], rank: int, + init_optimizer: bool = True, ) -> RuntimeConfig: """Validate configuration and prepare runtime settings. @@ -290,6 +330,10 @@ def validate_and_prepare_config( config: Policy configuration dictionary processor: Optional processor for multimodal models rank: Current process rank + init_optimizer: Whether the worker will train (and therefore needs an + optimizer). Trainable workers must keep their weights in float32 to + preserve master-weight precision; this flag is consulted to reject + a non-float32 ``automodel_cfg.load_precision`` on such workers. Returns: RuntimeConfig named tuple containing validated configuration values @@ -327,9 +371,20 @@ def validate_and_prepare_config( # Parse precision precision = config["precision"] - if precision not in STRING_TO_DTYPE: - raise ValueError(f"Unknown precision: {precision}") - dtype = STRING_TO_DTYPE[precision] + dtype = _precision_to_dtype(precision, "precision") + load_precision = config["automodel_cfg"].get("load_precision") + model_load_dtype = ( + _precision_to_dtype(load_precision, "automodel_cfg.load_precision") + if load_precision is not None + else None + ) + if init_optimizer and model_load_dtype not in (None, torch.float32): + raise ValueError( + "automodel_cfg.load_precision must be float32 for trainable workers " + "(init_optimizer=True) to preserve master-weight precision; it is " + "only intended for reference/teacher workers, but got " + f"{load_precision!r}." + ) # Get other configuration values cpu_offload = config["automodel_cfg"]["cpu_offload"] @@ -370,7 +425,7 @@ def validate_and_prepare_config( # Load model config model_config = AutoConfig.from_pretrained( model_name, - torch_dtype=dtype, + torch_dtype=model_load_dtype if model_load_dtype is not None else dtype, trust_remote_code=True, attn_implementation="flash_attention_2" if enable_seq_packing else None, **hf_config_overrides, @@ -440,6 +495,7 @@ def validate_and_prepare_config( model_class=model_class, model_config=model_config, hf_config_overrides=hf_config_overrides, + model_load_dtype=model_load_dtype, allow_flash_attn_args=allow_flash_attn_args, attn_impl=attn_impl, dtype=dtype, @@ -779,6 +835,9 @@ def setup_model_and_optimizer( # optimizer, so there are no master weights to protect. load_dtype = runtime_config.dtype + if runtime_config.model_load_dtype is not None: + load_dtype = runtime_config.model_load_dtype + # Validate CP configuration with model type before from_pretrained if cp_size > 1: if model_config.model_type == "gemma3": @@ -932,27 +991,50 @@ def setup_model_and_optimizer( activation_checkpointing=config["automodel_cfg"]["activation_checkpointing"], ) + # Automodel's NemotronH parallelization strategy crashes (and would silently + # skip Mamba CP) on the HF remote-code model at context_parallel_size > 1. + # Must run per worker process, before from_pretrained parallelizes the model. + # Gated on CP > 1 so a CP=1 run keeps byte-identical model construction. + if config["automodel_cfg"]["context_parallel_size"] > 1: + apply_nemotron_h_cp_fixes() + + shard_before_load = bool(config["automodel_cfg"].get("shard_before_load")) + print( + f"[Rank {rank}] Initializing model via from_pretrained " + f"(load_dtype={load_dtype}, precision={runtime_config.dtype}, " + f"shard_before_load={shard_before_load})..." + ) + # Create model via from_pretrained - handles meta device init, parallelization, # LoRA, and base weight loading internally - model = model_class.from_pretrained( - model_name, - distributed_setup=distributed_setup, - peft_config=peft_config, - attn_implementation=attn_impl, - torch_dtype=load_dtype, - trust_remote_code=True, - sdpa_method=sdpa_method, - **from_pretrained_kwargs, - **automodel_kwargs, - ) + with _force_shard_before_load(shard_before_load): + model = model_class.from_pretrained( + model_name, + distributed_setup=distributed_setup, + peft_config=peft_config, + attn_implementation=attn_impl, + torch_dtype=load_dtype, + trust_remote_code=True, + sdpa_method=sdpa_method, + **from_pretrained_kwargs, + **automodel_kwargs, + ) print(model) # Compute model metadata after from_pretrained model_state_dict_keys = list(model.state_dict().keys()) is_moe_model = any(["expert" in key for key in model_state_dict_keys]) + # force_hf means the HF implementation was loaded even though the architecture + # has a custom entry in ModelRegistry, so the registry lookup alone would + # misreport it as non-HF. That matters for the autocast gate below: an HF MoE + # model needs autocast to reconcile the fp32 activations produced by + # MixedPrecisionPolicy(output_dtype=float32) with its bf16 params, and without + # it the first unwrapped matmul raises "expected mat1 and mat2 to have the same + # dtype, float != c10::BFloat16" at lm_head. is_hf_model = ( - model_config.architectures[0] not in ModelRegistry.model_arch_name_to_cls + automodel_kwargs.get("force_hf", False) + or model_config.architectures[0] not in ModelRegistry.model_arch_name_to_cls ) # Autocast is disabled for custom MoE models (non-HF) to avoid numerical issues autocast_enabled = not (is_moe_model and not is_hf_model) @@ -969,6 +1051,28 @@ def setup_model_and_optimizer( if is_tied_lm_head: model.tie_weights() + # Freeze parameters matching the configured substrings. A submodule that never + # receives gradients (e.g. the nemotron_h `mtp.*` head during distillation, where + # loss flows only through lm_head) still has requires_grad=True by default, so the + # optimizer expects state for it on resume but none was saved -> "Missing key in + # checkpoint state_dict: optim.state.mtp...". Setting requires_grad=False makes + # save/load symmetric. Default [] -> no-op. + freeze_patterns = config.get("freeze_parameter_patterns") or [] + if freeze_patterns: + n_matched = 0 + n_frozen = 0 + for name, param in model.named_parameters(): + if any(pat in name for pat in freeze_patterns): + n_matched += 1 + if param.requires_grad: + param.requires_grad_(False) + n_frozen += 1 + if rank == 0: + print( + f"Freeze patterns matched {n_matched} parameters; newly froze " + f"{n_frozen}: freeze_parameter_patterns={list(freeze_patterns)}" + ) + # CPU offload if needed if cpu_offload: # Move buffers to CPU for FSDP modules diff --git a/nemo_rl/models/automodel/train.py b/nemo_rl/models/automodel/train.py index 8feeb008475..77d0aeab386 100644 --- a/nemo_rl/models/automodel/train.py +++ b/nemo_rl/models/automodel/train.py @@ -48,6 +48,11 @@ need_top_k_or_top_p_filtering, ) from nemo_rl.algorithms.loss import SequencePackingLossWrapper, prepare_loss_input +from nemo_rl.algorithms.loss.csd import CSDLossFn +from nemo_rl.algorithms.loss.csd_packing import ( + CSD_LOGICAL_METRICS_KEY, + CSDSequencePackingLossWrapper, +) from nemo_rl.algorithms.loss.interfaces import LossFunction, LossInputType from nemo_rl.algorithms.utils import mask_filtered_logprobs_outside_tokens from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -74,6 +79,7 @@ "LossPostProcessor", "LogprobsPostProcessor", "TopkLogitsPostProcessor", + "SparseLogitsPostProcessor", "FullLogitsPostProcessor", "ScorePostProcessor", ] @@ -107,13 +113,35 @@ def _build_model_batch( if processed_inputs.position_ids is not None: model_batch["position_ids"] = processed_inputs.position_ids if processed_inputs.has_flash_attention: - model_batch["flash_attn_kwargs"] = processed_inputs.flash_attn_kwargs + flash_kwargs = processed_inputs.flash_attn_kwargs + if hasattr(flash_kwargs, "cu_seq_lens_q"): + # Current Transformers/Automodel forwards these kwargs directly to + # flash-attention. Keep the legacy nested object for generic + # callers, but use the verified upstream spelling for planned + # packing so raw segment boundaries isolate logical samples. + model_batch.update( + { + "cu_seq_lens_q": flash_kwargs.cu_seq_lens_q, + "cu_seq_lens_k": flash_kwargs.cu_seq_lens_k, + "max_length_q": flash_kwargs.max_length_q, + "max_length_k": flash_kwargs.max_length_k, + } + ) + else: + model_batch["flash_attn_kwargs"] = flash_kwargs if processed_inputs.is_multimodal: model_batch.update( filter_multimodal_kwargs_for_model(model, processed_inputs.vlm_kwargs) ) - model_batch.pop("flash_attn_kwargs", None) + for key in ( + "flash_attn_kwargs", + "cu_seq_lens_q", + "cu_seq_lens_k", + "max_length_q", + "max_length_k", + ): + model_batch.pop(key, None) is_gemma3 = isinstance(model, Gemma3ForCausalLM) or isinstance( model, Gemma3ForConditionalGeneration @@ -128,7 +156,14 @@ def _build_model_batch( ) if is_reward_model or not allow_flash_attn_args: - model_batch.pop("flash_attn_kwargs", None) + for key in ( + "flash_attn_kwargs", + "cu_seq_lens_q", + "cu_seq_lens_k", + "max_length_q", + "max_length_k", + ): + model_batch.pop(key, None) drop_explicit_only_kwargs_not_in_forward(model, model_batch) @@ -304,6 +339,8 @@ def forward_with_post_processing_fn( processed_mb: ProcessedMicrobatch, global_valid_seqs: Optional[torch.Tensor] = None, global_valid_toks: Optional[torch.Tensor] = None, + global_valid_kd_toks: Optional[torch.Tensor] = None, + global_valid_chunks_by_idx: Optional[dict[int, torch.Tensor]] = None, sampling_params: Optional[TrainingSamplingParams] = None, sequence_dim: int = 1, ) -> Tuple[Any, dict[str, Any], ProcessedMicrobatch]: @@ -323,6 +360,7 @@ def forward_with_post_processing_fn( processed_mb: Pre-fetched ProcessedMicrobatch containing data and processed inputs global_valid_seqs: Global valid sequence count for loss normalization global_valid_toks: Global valid token count for loss normalization + global_valid_chunks_by_idx: Global valid alignment-chunk counts by teacher sampling_params: Sampling parameters (top-k, top-p, temperature) sequence_dim: Sequence dimension @@ -361,6 +399,7 @@ def forward_with_post_processing_fn( LossPostProcessor, LogprobsPostProcessor, TopkLogitsPostProcessor, + SparseLogitsPostProcessor, FullLogitsPostProcessor, ), ): @@ -378,11 +417,13 @@ def forward_with_post_processing_fn( global_valid_seqs=global_valid_seqs, global_valid_toks=global_valid_toks, cp_sharder=cp_sharder, + global_valid_kd_toks=global_valid_kd_toks, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, sequence_dim=sequence_dim, ) elif isinstance( post_processing_fn, - (LogprobsPostProcessor, TopkLogitsPostProcessor), + (LogprobsPostProcessor, TopkLogitsPostProcessor, SparseLogitsPostProcessor), ): result = post_processing_fn( logits=logits, @@ -395,9 +436,17 @@ def forward_with_post_processing_fn( ) if isinstance(post_processing_fn, LogprobsPostProcessor): metrics = {"logprobs": result} - else: + elif isinstance(post_processing_fn, TopkLogitsPostProcessor): vals, idx = result metrics = {"topk_logits": vals, "topk_indices": idx} + else: + vals, idx, log_z, gt_in_topk = result + metrics = { + "topk_logits": vals, + "topk_indices": idx, + "log_z": log_z, + "gt_in_topk": gt_in_topk, + } elif isinstance(post_processing_fn, FullLogitsPostProcessor): result = post_processing_fn( logits=logits, @@ -433,6 +482,8 @@ def automodel_forward_backward( allow_flash_attn_args: bool = True, global_valid_seqs: Optional[torch.Tensor] = None, global_valid_toks: Optional[torch.Tensor] = None, + global_valid_kd_toks: Optional[torch.Tensor] = None, + global_valid_chunks_by_idx: Optional[dict[int, torch.Tensor]] = None, sampling_params: Optional[TrainingSamplingParams] = None, sequence_dim: int = 1, dp_size: int = 1, @@ -462,6 +513,7 @@ def automodel_forward_backward( allow_flash_attn_args: Whether to pass flash_attn_kwargs to model global_valid_seqs: Global valid sequence count for loss normalization global_valid_toks: Global valid token count for loss normalization + global_valid_chunks_by_idx: Global valid alignment-chunk counts by teacher sampling_params: Sampling parameters (top-k, top-p, temperature) sequence_dim: Sequence dimension dp_size: Data parallel size @@ -509,6 +561,8 @@ def automodel_forward_backward( processed_mb=processed_mb, global_valid_seqs=global_valid_seqs, global_valid_toks=global_valid_toks, + global_valid_kd_toks=global_valid_kd_toks, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, sampling_params=sampling_params, sequence_dim=sequence_dim, ) @@ -524,11 +578,19 @@ def automodel_forward_backward( if not is_dummy: ## scale by the number of global batches so we get the correct ## value when summing metrics across all microbatches - for k in metrics.keys(): - if "_min" in k or "_max" in k: - continue - - metrics[k] /= num_global_batches + logical_metrics = metrics.get(CSD_LOGICAL_METRICS_KEY) + if logical_metrics is not None: + for logical_metric in logical_metrics: + for key in logical_metric: + if "_min" not in key and "_max" not in key: + logical_metric[key] /= num_global_batches + for key in metrics: + if ( + key != CSD_LOGICAL_METRICS_KEY + and "_min" not in key + and "_max" not in key + ): + metrics[key] /= num_global_batches else: # Zero out loss for dummy batches result = result * 0 @@ -616,6 +678,8 @@ def __call__( global_valid_toks: torch.Tensor, *, cp_sharder: Optional[ContextParallelSharder], + global_valid_kd_toks: Optional[torch.Tensor] = None, + global_valid_chunks_by_idx: Optional[dict[int, torch.Tensor]] = None, sequence_dim: int = 1, ) -> tuple[torch.Tensor, dict[str, Any]]: """Compute loss from logits. @@ -628,6 +692,9 @@ def __call__( global_valid_toks: Global valid token count cp_sharder: Per-microbatch Automodel sequence-layout owner, or None when context parallelism is inactive. + global_valid_kd_toks: Global valid knowledge-distillation token count + for losses normalized over aligned teacher/student tokens. + global_valid_chunks_by_idx: Global valid alignment-chunk counts by teacher sequence_dim: Sequence dimension Returns: @@ -662,31 +729,52 @@ def __call__( ), cp_sharder=token_layout, ) + extra_loss_kwargs: dict[str, Any] = {} + if isinstance(self.loss_fn, CSDLossFn) and global_valid_chunks_by_idx: + extra_loss_kwargs["global_valid_chunks_by_idx"] = global_valid_chunks_by_idx # Wrap loss function for sequence packing if needed if self.enable_seq_packing: - cu_seqlens_q_cpu = to_cpu_int_tuple( - processed_inputs.flash_attn_kwargs.cu_seqlens_q + is_csd_loss = isinstance(self.loss_fn, CSDLossFn) + packing_wrapper = ( + CSDSequencePackingLossWrapper + if is_csd_loss + else SequencePackingLossWrapper ) - loss_fn = SequencePackingLossWrapper( + loss_fn = packing_wrapper( loss_fn=self.loss_fn, prepare_fn=prepare_loss_input_wrapped, - cu_seqlens_q=cu_seqlens_q_cpu, - cu_seqlens_q_padded=cu_seqlens_q_cpu, + cu_seqlens_q=to_cpu_int_tuple( + processed_inputs.flash_attn_kwargs.cu_seqlens_q + ), + cu_seqlens_q_padded=to_cpu_int_tuple( + getattr( + processed_inputs.flash_attn_kwargs, + "cu_seqlens_q_padded", + processed_inputs.flash_attn_kwargs.cu_seqlens_q, + ) + ), ) + if is_csd_loss: + extra_loss_kwargs["global_valid_kd_toks"] = global_valid_kd_toks loss, loss_metrics = loss_fn( logits, data_dict, global_valid_seqs, global_valid_toks, + **extra_loss_kwargs, ) else: loss_input, data_dict = prepare_loss_input_wrapped( logits, data_dict, self.loss_fn ) + direct_loss_kwargs = dict(extra_loss_kwargs) + if isinstance(self.loss_fn, CSDLossFn): + direct_loss_kwargs["global_valid_kd_toks"] = global_valid_kd_toks loss, loss_metrics = self.loss_fn( data=data_dict, global_valid_seqs=global_valid_seqs, global_valid_toks=global_valid_toks, + **direct_loss_kwargs, **loss_input, ) @@ -787,13 +875,21 @@ def __call__( dtype=token_logprobs.dtype, device=token_logprobs.device, ) - cu_seqlens_cpu = to_cpu_int_tuple( + cu_seqlens = to_cpu_int_tuple( processed_inputs.flash_attn_kwargs.cu_seqlens_q ) + cu_seqlens_padded = to_cpu_int_tuple( + getattr( + processed_inputs.flash_attn_kwargs, + "cu_seqlens_q_padded", + processed_inputs.flash_attn_kwargs.cu_seqlens_q, + ) + ) input_lengths_cpu = to_cpu_int_tuple(input_lengths) for i in range(original_batch_size): - start = cu_seqlens_cpu[i] + 1 - end = cu_seqlens_cpu[i + 1] + raw_length = cu_seqlens[i + 1] - cu_seqlens[i] + start = cu_seqlens_padded[i] + 1 + end = cu_seqlens_padded[i] + raw_length seq_len_actual = input_lengths_cpu[i] unpacked_logprobs[i, 1:seq_len_actual] = token_logprobs[0, start:end] token_logprobs = unpacked_logprobs @@ -999,14 +1095,22 @@ def __call__( device=idx.device, ) - cu_seqlens_cpu = to_cpu_int_tuple( + cu_seqlens = to_cpu_int_tuple( processed_inputs.flash_attn_kwargs.cu_seqlens_q ) + cu_seqlens_padded = to_cpu_int_tuple( + getattr( + processed_inputs.flash_attn_kwargs, + "cu_seqlens_q_padded", + processed_inputs.flash_attn_kwargs.cu_seqlens_q, + ) + ) input_lengths_cpu = to_cpu_int_tuple(input_lengths) for i in range(original_batch_size): - start = cu_seqlens_cpu[i] - end = cu_seqlens_cpu[i + 1] + raw_length = cu_seqlens[i + 1] - cu_seqlens[i] + start = cu_seqlens_padded[i] + end = start + raw_length seq_len_actual = input_lengths_cpu[i] # Extract the corresponding portion from packed results @@ -1020,6 +1124,252 @@ def __call__( return vals, idx +class SparseLogitsPostProcessor: + """Compute sparse teacher payloads over TP and contiguous CP windows.""" + + def __init__( + self, + cfg: PolicyConfig, + tp_mesh: Any, + cp_mesh: Any, + cp_size: int, + k: int, + temperature: float, + vocab_size: Optional[int], + gt_filter_topk: int = 0, + ) -> None: + self.cfg = cfg + self.tp_mesh = tp_mesh + self.cp_mesh = cp_mesh + self.cp_size = cp_size + self.k = k + self.temperature = temperature + self.vocab_size = vocab_size + self.gt_filter_topk = gt_filter_topk + + @staticmethod + def _normalize_force_ids( + logits: torch.Tensor, + force_ids: Optional[torch.Tensor], + ) -> Optional[torch.Tensor]: + if force_ids is None: + return None + force_ids = force_ids.to(device=logits.device, dtype=torch.long) + if force_ids.shape[:2] == logits.shape[:2]: + return force_ids + common_batch = min(int(force_ids.shape[0]), int(logits.shape[0])) + common_seq = min(int(force_ids.shape[1]), int(logits.shape[1])) + if common_batch <= 0 or common_seq <= 0: + return None + trimmed = torch.full( + logits.shape[:2], -1, device=logits.device, dtype=torch.long + ) + trimmed[:common_batch, :common_seq] = force_ids[:common_batch, :common_seq] + return trimmed + + def __call__( + self, + logits: torch.Tensor, + data_dict: BatchedDataDict[Any], + processed_inputs: ProcessedInputs, + original_batch_size: int, + original_seq_len: int, + *, + cp_sharder: Optional[ContextParallelSharder], + sequence_dim: int = 1, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + del processed_inputs, original_batch_size, original_seq_len + local_logits = ( + logits.to_local() if isinstance(logits, DTensor) else logits + ).to(torch.float32) + tp_group = self.tp_mesh.get_group() + tp_rank = torch.distributed.get_rank(tp_group) + tp_size = torch.distributed.get_world_size(tp_group) + local_vocab_size = int(local_logits.shape[-1]) + if self.vocab_size is None and not isinstance(logits, DTensor): + raise ValueError( + "vocab_size is required for raw tensor logits under tensor parallelism" + ) + full_vocab_size = int(self.vocab_size or local_vocab_size * tp_size) + if full_vocab_size <= 0 or full_vocab_size > local_vocab_size * tp_size: + raise ValueError( + "vocab_size must be within the physical distributed lm-head " + f"width; got vocab_size={full_vocab_size}, " + f"physical_width={local_vocab_size * tp_size}." + ) + vocab_is_replicated = tp_size > 1 and local_vocab_size == full_vocab_size + vocab_start = ( + tp_rank * full_vocab_size + if vocab_is_replicated + else tp_rank * local_vocab_size + ) + valid_local_vocab = max(0, min(local_vocab_size, full_vocab_size - vocab_start)) + if valid_local_vocab < local_vocab_size: + local_logits = local_logits.clone() + local_logits[..., valid_local_vocab:] = -torch.inf + + actual_k = min(int(self.k), full_vocab_size) + real_logits: Optional[torch.Tensor] = None + if tp_size == 1: + # Match the reference TP1 producer exactly: slice padded lm-head + # columns before both reductions, then use native logsumexp/topk. + real_logits = local_logits[..., :full_vocab_size] + log_z = torch.logsumexp(real_logits / float(self.temperature), dim=-1) + top_vals, top_idx = torch.topk(real_logits, k=actual_k, dim=-1) + else: + scaled_logits = local_logits / float(self.temperature) + global_max = scaled_logits.amax(dim=-1) + torch.distributed.all_reduce( + global_max, op=torch.distributed.ReduceOp.MAX, group=tp_group + ) + exp_sum = torch.exp(scaled_logits - global_max.unsqueeze(-1)).sum(dim=-1) + torch.distributed.all_reduce( + exp_sum, op=torch.distributed.ReduceOp.SUM, group=tp_group + ) + log_z = global_max + torch.log(exp_sum) + top_vals, top_idx = distributed_vocab_topk( + local_logits, + k=actual_k, + tp_group=tp_group, + vocab_start_index=vocab_start, + vocab_end_index=vocab_start + local_vocab_size, + ) + + force_ids = self._normalize_force_ids( + local_logits, + self._localize_force_ids( + data_dict.get("force_include_token_ids"), + local_logits.shape[1], + cp_sharder=cp_sharder, + ), + ) + gt_in_topk: Optional[torch.Tensor] = None + if self.gt_filter_topk > 0: + if force_ids is None: + raise ValueError( + "gt_filter_topk requires force_include_token_ids in teacher data" + ) + filter_k = min(self.gt_filter_topk, full_vocab_size) + if tp_size == 1: + assert real_logits is not None + filter_idx = torch.topk(real_logits, k=filter_k, dim=-1).indices + else: + _, filter_idx = distributed_vocab_topk( + local_logits, + k=filter_k, + tp_group=tp_group, + vocab_start_index=vocab_start, + vocab_end_index=vocab_start + local_vocab_size, + ) + valid_force = (force_ids >= 0) & (force_ids < full_vocab_size) + gt_in_topk = valid_force & ( + filter_idx.to(torch.long) == force_ids.unsqueeze(-1) + ).any(dim=-1) + + if force_ids is not None and top_vals.shape[-1] > 0: + valid_force = (force_ids >= 0) & (force_ids < full_vocab_size) + already_present = (top_idx.to(torch.long) == force_ids.unsqueeze(-1)).any( + dim=-1 + ) + need_force = valid_force & ~already_present + if tp_size == 1: + assert real_logits is not None + forced_vals = real_logits.gather( + dim=-1, + index=force_ids.clamp(0, full_vocab_size - 1).unsqueeze(-1), + ).squeeze(-1) + else: + owned = (force_ids >= vocab_start) & ( + force_ids < vocab_start + valid_local_vocab + ) + local_force_idx = (force_ids - vocab_start).clamp( + min=0, max=max(local_vocab_size - 1, 0) + ) + forced_vals = local_logits.gather( + dim=-1, index=local_force_idx.unsqueeze(-1) + ).squeeze(-1) + forced_vals = torch.where( + owned, forced_vals, torch.full_like(forced_vals, -torch.inf) + ) + torch.distributed.all_reduce( + forced_vals, op=torch.distributed.ReduceOp.MAX, group=tp_group + ) + replace_slot = top_vals.argmin(dim=-1) + replacement_vals = torch.where( + need_force, + forced_vals, + top_vals.gather(-1, replace_slot.unsqueeze(-1)).squeeze(-1), + ) + replacement_idx = torch.where( + need_force, + force_ids, + top_idx.to(torch.long) + .gather(-1, replace_slot.unsqueeze(-1)) + .squeeze(-1), + ) + top_vals = top_vals.clone() + top_idx = top_idx.clone() + top_vals.scatter_( + -1, replace_slot.unsqueeze(-1), replacement_vals.unsqueeze(-1) + ) + top_idx.scatter_( + -1, + replace_slot.unsqueeze(-1), + replacement_idx.to(top_idx.dtype).unsqueeze(-1), + ) + + sorted_idx, sort_perm = torch.sort(top_idx.to(torch.int32), dim=-1) + sorted_vals = torch.gather(top_vals, dim=-1, index=sort_perm) + if cp_sharder is not None: + sorted_vals = self._gather_and_slice_cp( + sorted_vals, cp_sharder, sequence_dim + ) + sorted_idx = self._gather_and_slice_cp(sorted_idx, cp_sharder, sequence_dim) + log_z = self._gather_and_slice_cp(log_z, cp_sharder, sequence_dim) + if gt_in_topk is not None: + gt_in_topk = self._gather_and_slice_cp( + gt_in_topk, cp_sharder, sequence_dim + ) + return sorted_vals, sorted_idx, log_z, gt_in_topk + + def _gather_and_slice_cp( + self, + tensor: torch.Tensor, + cp_sharder: ContextParallelSharder, + sequence_dim: int, + ) -> torch.Tensor: + """Restore canonical order, then emit this CP rank's IPC window.""" + full = cp_sharder.gather_token_tensor(tensor, seq_dim=sequence_dim, trim=True) + full_seq_len = int(full.shape[sequence_dim]) + if full_seq_len % self.cp_size != 0: + raise ValueError( + "CSD teacher sequence length must be divisible by the " + "teacher context parallel size, but got " + f"sequence_length={full_seq_len}, cp_size={self.cp_size}." + ) + cp_rank = torch.distributed.get_rank(self.cp_mesh.get_group()) + local_len = full_seq_len // self.cp_size + return full.narrow(sequence_dim, cp_rank * local_len, local_len).contiguous() + + def _localize_force_ids( + self, + force_ids: Optional[torch.Tensor], + local_seq_len: int, + *, + cp_sharder: Optional[ContextParallelSharder], + ) -> Optional[torch.Tensor]: + """Map canonical labels into the model-owned local CP layout.""" + if force_ids is None or cp_sharder is None: + return force_ids + local = cp_sharder.shard_token_tensor(force_ids, seq_dim=1, fill=-1) + if local.shape[1] != local_seq_len: + raise ValueError( + "CP-local force-id length does not match local teacher logits: " + f"force_ids={local.shape[1]}, logits={local_seq_len}." + ) + return local + + class FullLogitsPostProcessor: """Export this rank's raw teacher logits (full vocab, no reduction at the worker). @@ -1029,7 +1379,9 @@ class FullLogitsPostProcessor: are supported and may differ from the student's: under TP each rank emits its vocab shard, under CP it allgathers and re-emits its contiguous seq slice; the IPC consumer reassembles the global ``[B, T_t, V_t]``. - Sequence packing raises ``NotImplementedError``. + With CP1 sequence packing, packed logits are restored to one dense logical + row per input sample before export. Sparse packed IPC remains a separate, + guarded path. """ def __init__( @@ -1065,22 +1417,67 @@ def __call__( cp_sharder: Optional[ContextParallelSharder], sequence_dim: int = 1, ) -> torch.Tensor: - if self.enable_seq_packing: - raise NotImplementedError( - "FullLogitsPostProcessor: sequence packing is not supported in v0." - ) if isinstance(logits, DTensor): logits = logits.to_local() # fp32 for the consumer's precision-sensitive log_softmax / top-k / # projection (KL math), and a dtype-consistent IPC buffer producer<->consumer. logits = logits.to(torch.float32) + if self.enable_seq_packing: + if self.cp_size != 1: + raise NotImplementedError( + "Packed Automodel full-logit export currently supports only " + "context_parallel_size=1." + ) + if logits.ndim != 3 or logits.shape[0] != 1: + raise ValueError( + "Packed Automodel logits must have shape [1, T, V], got " + f"{tuple(logits.shape)}." + ) + cu_seqlens = processed_inputs.flash_attn_kwargs.cu_seqlens_q + cu_seqlens_padded = getattr( + processed_inputs.flash_attn_kwargs, + "cu_seqlens_q_padded", + cu_seqlens, + ) + if int(cu_seqlens.numel()) != original_batch_size + 1: + raise ValueError( + "Packed Automodel cu_seqlens cardinality does not match the " + f"logical batch: entries={cu_seqlens.numel()}, " + f"batch={original_batch_size}." + ) + unpacked = logits.new_zeros( + (original_batch_size, original_seq_len, logits.shape[-1]) + ) + input_lengths = data_dict["input_lengths"] + for sample_index in range(original_batch_size): + raw_start = int(cu_seqlens[sample_index].item()) + raw_end = int(cu_seqlens[sample_index + 1].item()) + raw_length = int(input_lengths[sample_index].item()) + if raw_end - raw_start != raw_length: + raise ValueError( + "Packed Automodel raw boundary disagrees with input_lengths " + f"for sample {sample_index}: cu={raw_end - raw_start}, " + f"input_length={raw_length}." + ) + padded_start = int(cu_seqlens_padded[sample_index].item()) + physical_end = padded_start + raw_length + if raw_length > original_seq_len or physical_end > logits.shape[1]: + raise ValueError( + "Packed Automodel logits cannot be restored for sample " + f"{sample_index}: physical_end={physical_end}, raw_length={raw_length}, " + f"packed_width={logits.shape[1]}, logical_width={original_seq_len}." + ) + unpacked[sample_index, :raw_length].copy_( + logits[0, padded_start:physical_end] + ) + return unpacked + # Automodel's CP layout is not contiguous per rank, but the IPC consumer - # routes by contiguous ``global_seq_start`` over the teacher CP group. - # Restore canonical order (trimming Automodel's CP padding) and emit this - # rank's contiguous slice, else heterogeneous teacher_cp != student_cp - # lands teacher data at the wrong seq positions in the consumer's dest - # tensor. + # routes by contiguous ``global_seq_start`` over the teacher CP + # group. Restore global order and emit this rank's contiguous slice, else + # heterogeneous teacher_cp != student_cp lands teacher data at the wrong + # seq positions in the consumer's dest tensor. if cp_sharder is not None and self.cp_mesh is not None: full = cp_sharder.gather_token_tensor( logits, seq_dim=sequence_dim, trim=True diff --git a/nemo_rl/models/megatron/data.py b/nemo_rl/models/megatron/data.py index 19a30bcade6..b2701366c8a 100644 --- a/nemo_rl/models/megatron/data.py +++ b/nemo_rl/models/megatron/data.py @@ -15,7 +15,7 @@ from contextlib import nullcontext from dataclasses import dataclass from math import lcm -from typing import Any, Iterator, Optional, Sequence, Tuple +from typing import Any, Iterable, Iterator, Optional, Sequence, Tuple import torch from megatron.bridge.training.utils.packed_seq_utils import ( @@ -29,6 +29,7 @@ ) from megatron.core.utils import StragglerDetector, get_batch_on_this_cp_rank +from nemo_rl.algorithms.loss.csd_packing import resolve_lockstep_bin_geometry from nemo_rl.algorithms.loss.interfaces import LossFunction, LossType from nemo_rl.data.megatron_sft_packed import ( MEGATRON_SFT_PACKED_BATCH_FIELDS, @@ -64,8 +65,8 @@ class ProcessedInputs: position_ids: Optional[torch.Tensor] packed_seq_params: Optional[PackedSeqParams] cu_seqlens_padded: Optional[torch.Tensor] - mtp_loss_mask: Optional[torch.Tensor] = None padding_mask: Optional[torch.Tensor] = None + mtp_loss_mask: Optional[torch.Tensor] = None routed_experts: Optional[torch.Tensor] = None routed_experts_cp_sharded: Optional[torch.Tensor] = None labels_cp_sharded: Optional[torch.Tensor] = None @@ -120,8 +121,8 @@ class ProcessedMicrobatch: position_ids: Optional[torch.Tensor] packed_seq_params: Optional[PackedSeqParams] cu_seqlens_padded: Optional[torch.Tensor] - mtp_loss_mask: Optional[torch.Tensor] = None padding_mask: Optional[torch.Tensor] = None + mtp_loss_mask: Optional[torch.Tensor] = None routed_experts: Optional[torch.Tensor] = None routed_experts_cp_sharded: Optional[torch.Tensor] = None labels_cp_sharded: Optional[torch.Tensor] = None @@ -298,10 +299,10 @@ def make_processed_microbatch_iterator( position_ids=processed_inputs.position_ids, packed_seq_params=processed_inputs.packed_seq_params, cu_seqlens_padded=processed_inputs.cu_seqlens_padded, - mtp_loss_mask=processed_inputs.mtp_loss_mask, padding_mask=processed_inputs.padding_mask, labels_cp_sharded=processed_inputs.labels_cp_sharded, loss_mask_cp_sharded=processed_inputs.loss_mask_cp_sharded, + mtp_loss_mask=processed_inputs.mtp_loss_mask, routed_experts=processed_inputs.routed_experts, routed_experts_cp_sharded=processed_inputs.routed_experts_cp_sharded, original_seq_length=processed_inputs.original_seq_length, @@ -365,6 +366,7 @@ def get_microbatch_iterator( mtp_enabled: bool = False, create_router_padding_mask: bool = False, allow_unmasked_chunkwise_cp: bool = False, + skip_keys: Optional[Iterable[str]] = None, ) -> Tuple[Iterator[ProcessedMicrobatch], int, int, int, int]: """Create a processed microbatch iterator from a batch of data. @@ -382,6 +384,9 @@ def get_microbatch_iterator( on supported MoE models, for packed and nonpacked inputs alike. allow_unmasked_chunkwise_cp: Suppress router masks for prepacked HybridEP inputs after setup warned about unsupported HybridBlock chunkwise CP. + skip_keys: Keys whose dim 1 is NOT the sequence axis and must be exempt + from the sequence-dim validation (e.g. cross-tokenizer ride-along + tensors). Forwarded to ``get_and_validate_seqlen``. Returns: Tuple containing the iterator and metadata @@ -415,7 +420,7 @@ def get_microbatch_iterator( "multiple of that value" ) else: - _, seq_dim_size = get_and_validate_seqlen(data) + _, seq_dim_size = get_and_validate_seqlen(data, skip_keys=skip_keys) # Auto-detect seq_length_key if not provided if seq_length_key is None and cfg["sequence_packing"]["enabled"]: @@ -813,10 +818,10 @@ def process_microbatch( seq_lengths = None # Will be set if using packed sequences cu_seqlens = None cu_seqlens_padded = None - mtp_loss_mask = None padding_mask = None labels_cp_sharded = None loss_mask_cp_sharded = None + mtp_loss_mask = None media_token_validity_mask = None if is_direct_packed_row(data_dict): @@ -876,6 +881,14 @@ def process_microbatch( # Get sequence lengths and context parallel size seq_lengths = data_dict[seq_length_key] + planned_geometry = resolve_lockstep_bin_geometry( + data_dict, input_lengths=seq_lengths + ) + planned_physical_tokens = ( + planned_geometry.physical_tokens + if planned_geometry is not None + else pad_full_seq_to + ) prepacked = "cu_seqlens" in data_dict if prepacked: @@ -925,6 +938,11 @@ def process_microbatch( "Router padding masks require real-token validity through " "model-owned sequence packing, which is not supported." ) + if planned_geometry is not None: + raise NotImplementedError( + "Plan-driven packing cannot delegate materialization to " + "the model." + ) has_mtp_loss_mask = "mtp_loss_mask" in data_dict assert not has_mtp_loss_mask or delegate_mtp_loss_mask_to_model, ( "MTP training requires a self-packing VLM that advertises " @@ -1008,10 +1026,34 @@ def process_microbatch( seq_lengths, pad_individual_seqs_to_multiple_of, pad_packed_seq_to_multiple_of, - pad_full_seq_to, + planned_physical_tokens, cp_rank=get_context_parallel_rank(), cp_size=get_context_parallel_world_size(), ) + if planned_geometry is not None: + actual_raw_cu = tuple(int(value) for value in cu_seqlens.tolist()) + actual_padded_cu = tuple( + int(value) for value in cu_seqlens_padded.tolist() + ) + if actual_raw_cu != planned_geometry.raw_cu_seqlens: + raise ValueError( + "Megatron raw cumulative boundaries drifted from the " + f"selected lockstep bin: actual={actual_raw_cu}, " + f"planned={planned_geometry.raw_cu_seqlens}." + ) + if actual_padded_cu != planned_geometry.padded_cu_seqlens: + raise ValueError( + "Megatron padded cumulative boundaries drifted from " + f"the selected lockstep bin: actual={actual_padded_cu}, " + f"planned={planned_geometry.padded_cu_seqlens}." + ) + if input_ids.shape[1] != planned_geometry.physical_tokens: + raise ValueError( + "Megatron materialized physical size drifted from the " + f"selected lockstep bin: actual={input_ids.shape[1]}, " + f"planned={planned_geometry.physical_tokens}." + ) + if model_slices_context_parallel_inputs: packed_seq_params = PackedSeqParams( cu_seqlens_q=cu_seqlens, @@ -1082,6 +1124,32 @@ def process_microbatch( f"padding_mask shape {padding_mask.shape} must match " f"model input shape {input_ids_cp_sharded.shape}." ) + + # GPT's MoE router must exclude both physical padding and + # sample_mask=0 rows from aux/z-loss and capacity accounting. + # Models that own CP slicing (currently Nemotron Omni/LLaVA) + # do not accept a padding_mask forward kwarg, so preserve their + # existing packed multimodal contract until the wrapper plumbs + # the mask through to its language model. + if not model_slices_context_parallel_inputs: + router_padding_mask = _make_packed_router_padding_mask( + seq_lengths, + cu_seqlens_padded, + sample_mask=data_dict.get("sample_mask"), + cp_rank=get_context_parallel_rank(), + cp_size=get_context_parallel_world_size(), + device=input_ids_cp_sharded.device, + ) + padding_mask = ( + router_padding_mask + if padding_mask is None + else padding_mask | router_padding_mask + ) + assert padding_mask.shape == input_ids_cp_sharded.shape, ( + f"padding_mask shape {padding_mask.shape} must match " + f"model input shape {input_ids_cp_sharded.shape}." + ) + materialized_physical_tokens = int(input_ids.shape[1]) # routed_experts and the R3 trace token identity ride the SAME # per-seq zigzag CP sharding as input_ids, re-derived from # cu_seqlens_padded. @@ -1163,7 +1231,7 @@ def process_microbatch( seq_lengths, pad_individual_seqs_to_multiple_of, pad_packed_seq_to_multiple_of, - pad_full_seq_to, + materialized_physical_tokens, cp_rank=get_context_parallel_rank(), cp_size=get_context_parallel_world_size(), ) @@ -1202,7 +1270,7 @@ def process_microbatch( seq_lengths, pad_individual_seqs_to_multiple_of, pad_packed_seq_to_multiple_of, - pad_full_seq_to, + materialized_physical_tokens, cp_rank=get_context_parallel_rank(), cp_size=get_context_parallel_world_size(), ) @@ -1241,7 +1309,7 @@ def process_microbatch( seq_lengths, pad_individual_seqs_to_multiple_of, pad_packed_seq_to_multiple_of, - pad_full_seq_to, + materialized_physical_tokens, cp_rank=get_context_parallel_rank(), cp_size=get_context_parallel_world_size(), ) @@ -1282,7 +1350,57 @@ def process_microbatch( input_ids, data_dict["input_lengths"], ) - input_ids_cp_sharded = input_ids + cp_size = get_context_parallel_world_size() + attention_mask, _, position_ids = get_ltor_masks_and_position_ids( + data=input_ids, + eod_token=0, # used for loss_mask, which we don't use + pad_token=0, # used for loss_mask, which we don't use + reset_position_ids=False, + reset_attention_mask=False, + eod_mask_loss=False, + pad_mask_loss=False, + create_attention_mask=cp_size == 1, + ) + if cp_size > 1: + # Unpacked context parallelism. mcore expects each rank to be + # handed its load-balanced sequence shard -- chunks ``cp_rank`` + # and ``2*cp-1-cp_rank`` of ``2*cp`` -- exactly what + # Megatron-LM's get_batch_on_this_cp_rank produces and what the + # packed path above already does per sequence. position_ids ride + # the same layout. + # + # No dense [1, 1, S, S] mask is built here: it would be the + # plain causal mask (reset_position_ids / reset_attention_mask + # are both False), which the attention backend applies itself, + # and TE's context-parallel attention requires that implicit + # causal path rather than an arbitrary full-sequence mask. + if routed_experts is not None: + raise NotImplementedError( + "Router replay (routed_experts) with context parallelism " + "requires sequence packing, so routed_experts can ride " + "the same per-sequence CP layout as input_ids." + ) + assert input_ids.shape[1] % (cp_size * 2) == 0, ( + "Context parallelism load-balances the sequence into " + f"2*cp={cp_size * 2} chunks, so the padded sequence length " + f"({input_ids.shape[1]}) must be a multiple of it. Set " + "policy.make_sequence_length_divisible_by to " + "tensor_model_parallel_size * context_parallel_size * 2." + ) + cp_rank = get_context_parallel_rank() + input_ids_cp_sharded = _get_tokens_on_this_cp_rank( + input_ids, cp_rank, cp_size, seq_dim=1 + ) + position_ids = _get_tokens_on_this_cp_rank( + position_ids, cp_rank, cp_size, seq_dim=1 + ) + if padding_mask is not None: + padding_mask = _get_tokens_on_this_cp_rank( + padding_mask, cp_rank, cp_size, seq_dim=1 + ) + attention_mask = None + else: + input_ids_cp_sharded = input_ids verified_token_count = _verify_r3_trace_cp_token_alignment( source_input_ids=data_dict["input_ids"], source_routed_experts=data_dict.get("routed_experts"), @@ -1298,23 +1416,23 @@ def process_microbatch( cp_rank=get_context_parallel_rank(), cp_size=get_context_parallel_world_size(), ) - attention_mask, _, position_ids = get_ltor_masks_and_position_ids( - data=input_ids, - eod_token=0, # used for loss_mask, which we don't use - pad_token=0, # used for loss_mask, which we don't use - reset_position_ids=False, - reset_attention_mask=False, - eod_mask_loss=False, - pad_mask_loss=False, - ) if "mtp_loss_mask" in data_dict: mtp_loss_mask = data_dict["mtp_loss_mask"] - # Unpacked: rows still are samples, so the mask is already in the - # layout the model will see. + if cp_size > 1: + mtp_loss_mask = _get_tokens_on_this_cp_rank( + mtp_loss_mask, cp_rank, cp_size, seq_dim=1 + ) if "media_token_validity_mask" in data_dict: media_token_validity_mask = data_dict[ "media_token_validity_mask" ].bool() + if cp_size > 1: + media_token_validity_mask = _get_tokens_on_this_cp_rank( + media_token_validity_mask, + cp_rank, + cp_size, + seq_dim=1, + ) return ProcessedInputs( input_ids=input_ids, input_ids_cp_sharded=input_ids_cp_sharded, @@ -1322,10 +1440,10 @@ def process_microbatch( position_ids=position_ids, packed_seq_params=packed_seq_params, cu_seqlens_padded=cu_seqlens_padded, - mtp_loss_mask=mtp_loss_mask, padding_mask=padding_mask, labels_cp_sharded=labels_cp_sharded, loss_mask_cp_sharded=loss_mask_cp_sharded, + mtp_loss_mask=mtp_loss_mask, routed_experts=routed_experts, routed_experts_cp_sharded=routed_experts_cp_sharded, original_seq_length=original_seq_length, @@ -1376,6 +1494,51 @@ def _make_r3_trace_token_identity( ) +def _make_packed_router_padding_mask( + seq_lengths: torch.Tensor, + cu_seqlens_padded: torch.Tensor, + *, + sample_mask: Optional[torch.Tensor], + cp_rank: int, + cp_size: int, + device: torch.device, +) -> torch.Tensor: + """Build GPT's CP-local MoE padding mask in native per-sequence order.""" + if cu_seqlens_padded.numel() != seq_lengths.numel() + 1: + raise ValueError( + "Padded boundaries must contain one entry per packed sequence plus one." + ) + if sample_mask is not None and sample_mask.numel() != seq_lengths.numel(): + raise ValueError("sample_mask must contain one entry per packed sequence.") + + valid_segments: list[torch.Tensor] = [] + for sequence_index in range(int(seq_lengths.numel())): + raw_length = int(seq_lengths[sequence_index].item()) + padded_length = int( + ( + cu_seqlens_padded[sequence_index + 1] + - cu_seqlens_padded[sequence_index] + ).item() + ) + if padded_length < raw_length: + raise ValueError( + f"Packed sequence {sequence_index} has raw length {raw_length} " + f"but padded length {padded_length}." + ) + valid = torch.zeros(padded_length, dtype=torch.bool, device=device) + sample_is_valid = sample_mask is None or bool( + sample_mask[sequence_index].item() + ) + if sample_is_valid: + valid[:raw_length] = True + if cp_size > 1: + valid = _get_tokens_on_this_cp_rank( + valid, cp_rank=cp_rank, cp_size=cp_size, seq_dim=0 + ) + valid_segments.append(valid) + return ~torch.cat(valid_segments).unsqueeze(0) + + def _verify_r3_trace_cp_token_alignment( *, source_input_ids: torch.Tensor, @@ -1505,6 +1668,8 @@ def process_global_batch( - batch: The extracted batch - global_valid_seqs: Number of valid sequences across all ranks - global_valid_toks: Number of valid tokens across all ranks + - global_valid_chunks_by_idx: Number of valid cross-tokenizer alignment + chunks for each teacher across all data-parallel ranks """ batch = data.get_batch(batch_idx=batch_idx, batch_size=batch_size) @@ -1523,10 +1688,41 @@ def process_global_batch( local_valid_toks = torch.sum( batch["token_mask"][:, 1:] * batch["sample_mask"].unsqueeze(-1) ) + kd_token_mask = batch.get("kd_token_mask", batch.get("token_mask")) + if kd_token_mask is None: + local_valid_kd_toks = local_valid_toks + else: + local_valid_kd_toks = torch.sum( + kd_token_mask[:, 1:] * batch["sample_mask"].unsqueeze(-1) + ) + + local_valid_chunks_by_idx: dict[int, torch.Tensor] = {} + for key, pair_valid in batch.items(): + if not key.startswith("alignment_") or not key.endswith("_pair_valid"): + continue + teacher_idx = key.removeprefix("alignment_").removesuffix("_pair_valid") + if not teacher_idx.isdigit(): + continue + sample_mask = batch["sample_mask"].to(pair_valid.dtype).unsqueeze(-1) + local_valid_chunks_by_idx[int(teacher_idx)] = torch.sum( + pair_valid * sample_mask + ) - to_reduce = torch.tensor([local_valid_seqs, local_valid_toks]).cuda() + chunk_teacher_indices = sorted(local_valid_chunks_by_idx) + to_reduce = torch.stack( + [ + local_valid_seqs, + local_valid_toks, + local_valid_kd_toks, + *(local_valid_chunks_by_idx[i] for i in chunk_teacher_indices), + ] + ).cuda() torch.distributed.all_reduce(to_reduce, group=dp_group) - global_valid_seqs, global_valid_toks = to_reduce[0], to_reduce[1] + global_valid_seqs, global_valid_toks, global_valid_kd_toks = to_reduce[:3] + global_valid_chunks_by_idx = { + teacher_idx: to_reduce[offset + 3] + for offset, teacher_idx in enumerate(chunk_teacher_indices) + } if hasattr(loss_fn, "loss_type") and loss_fn.loss_type == LossType.TOKEN_LEVEL: assert "token_mask" in batch, ( @@ -1537,6 +1733,8 @@ def process_global_batch( "batch": batch, "global_valid_seqs": global_valid_seqs, "global_valid_toks": global_valid_toks, + "global_valid_kd_toks": global_valid_kd_toks, + "global_valid_chunks_by_idx": global_valid_chunks_by_idx, } @@ -1674,7 +1872,13 @@ def _pack_sequences_for_megatron( pad_packed_seq_to: Optional[int] = None, cp_rank: int = 0, cp_size: int = 1, -) -> tuple[torch.Tensor, PackedSeqParams, torch.Tensor, Optional[torch.Tensor]]: +) -> tuple[ + torch.Tensor, + torch.Tensor, + PackedSeqParams, + torch.Tensor, + torch.Tensor, +]: """Pack sequences for Megatron model processing with optional context parallelism. Args: @@ -1830,8 +2034,8 @@ def _pack_sequences_for_megatron( # total_tokens is required for PackedSeqParams.__post_init__ to build # seq_idx, which Mamba uses to reset SSM state at sample boundaries. packed_seq_params = PackedSeqParams( - cu_seqlens_q=cu_seqlens_padded, - cu_seqlens_kv=cu_seqlens_padded, + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, cu_seqlens_q_padded=cu_seqlens_padded, cu_seqlens_kv_padded=cu_seqlens_padded, max_seqlen_q=int(max_seqlen), @@ -2105,7 +2309,10 @@ def _unpack_sequences_from_megatron( return unpacked_output -def get_and_validate_seqlen(data: BatchedDataDict[Any]): +def get_and_validate_seqlen( + data: BatchedDataDict[Any], + skip_keys: Optional[Iterable[str]] = None, +): # dim 1 is always assumed to be the sequence dim, sanity check this here. # Skip multimodal fields: their dim 1 is num_images / num_patches, not # seqlen. Skip opd_full's teacher payload columns too: materialize() @@ -2114,10 +2321,14 @@ def get_and_validate_seqlen(data: BatchedDataDict[Any]): # each sequence's own real length regardless of how far the column is # padded) -- see codec.py's OPD_FULL_FIELDS exclusion. Their dim 1 can # therefore legitimately be shorter than every other column's. + # Additional cross-tokenizer ride-along fields are excluded by skip_keys. sequence_dim = 1 seq_dim_size = data["input_ids"].shape[sequence_dim] + skip = set(PACKED_MULTIMODAL_FIELDS) + if skip_keys is not None: + skip.update(skip_keys) for k, v in data.items(): - if k in PACKED_MULTIMODAL_FIELDS or k in OPD_FULL_FIELDS: + if k in skip or k in OPD_FULL_FIELDS: continue if torch.is_tensor(v) and not v.is_nested and len(v.shape) > 1: assert v.shape[sequence_dim] == seq_dim_size, ( diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 776f77f0fa3..d3ceed55d68 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -60,6 +60,7 @@ ) from megatron.bridge.training.state import GlobalState, TrainState from megatron.bridge.training.tokenizers.tokenizer import build_tokenizer +from megatron.bridge.training.utils.checkpoint_utils import is_hf_checkpoint_dir from megatron.bridge.training.utils.pg_utils import get_pg_collection from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size @@ -627,11 +628,17 @@ def validate_and_set_config( skip_weight_load=skip_weight_load, ) - final_padded_vocab_size = calculate_padded_vocab_size( - megatron_cfg.model.vocab_size, - megatron_cfg.model.make_vocab_size_divisible_by, - config["megatron_cfg"]["tensor_model_parallel_size"], - ) + model_vocab_size = int(megatron_cfg.model.vocab_size) + if megatron_cfg.model.should_pad_vocab: + final_padded_vocab_size = calculate_padded_vocab_size( + model_vocab_size, + megatron_cfg.model.make_vocab_size_divisible_by, + megatron_cfg.model.tensor_model_parallel_size, + ) + else: + # Match ModelProvider.provide(): an unpadded output head uses the raw + # configured vocabulary even when a larger divisible size exists. + final_padded_vocab_size = model_vocab_size return RuntimeConfig( megatron_cfg, @@ -659,6 +666,17 @@ def _get_hf_config_overrides_hash(overrides: dict[str, Any]) -> str: return hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:12] +def _use_megatron_fsdp(config: PolicyConfig) -> bool: + """Return whether the Megatron-Core FSDP data-parallel path is enabled. + + ``use_custom_fsdp`` was the name exposed by older NeMo RL configs, but + Megatron-Core has renamed that flag to ``use_megatron_fsdp``. Honor the + legacy spelling as a fallback for existing recipes. + """ + ddp_cfg = config.get("megatron_cfg", {}).get("distributed_data_parallel_config", {}) + return bool(ddp_cfg.get("use_megatron_fsdp", ddp_cfg.get("use_custom_fsdp", False))) + + def _resolve_iter_dir_from_root( path: str, not_found_msg: str, @@ -957,6 +975,23 @@ def validate_model_paths(config: PolicyConfig) -> tuple[str, str, bool]: # Existing HF path: cfg["model_name"] is an HF model name or local HF checkpoint. hf_model_name = config["model_name"] + + # Megatron-FSDP parameters are DTensors and cannot be initialized from the + # ordinary cached torch_dist conversion. Bridge supports loading HF weights + # directly into the FSDP shards, so resolve the local HF snapshot and use it + # as the pretrained source instead of creating/reusing a torch_dist cache. + if _use_megatron_fsdp(config): + if os.path.isdir(hf_model_name): + hf_snapshot_path = hf_model_name + else: + from huggingface_hub import snapshot_download + + hf_snapshot_path = snapshot_download( + repo_id=hf_model_name, + local_files_only=os.getenv("HF_HUB_OFFLINE") == "1", + ) + return hf_model_name, hf_snapshot_path, True + hf_config_overrides = config.get("hf_config_overrides", {}) or {} hf_model_subdir = hf_model_name @@ -1007,13 +1042,23 @@ def setup_model_config( """ pretrained_ckpt = config.get("pretrained_checkpoint") fmt = pretrained_ckpt["format"] if pretrained_ckpt is not None else None + use_megatron_fsdp = _use_megatron_fsdp(config) validate_router_replay_config(config) - derive_provider_from_hf = fmt == "megatron_lm" or ( - skip_weight_load and fmt != "megatron_bridge" + derive_provider_from_hf = ( + fmt == "megatron_lm" + or (fmt is None and use_megatron_fsdp) + or (skip_weight_load and fmt != "megatron_bridge") ) if derive_provider_from_hf: + # For megatron_lm format: build the model config from the HF architecture. + # pretrained_path has already been resolved to a specific iter dir by + # validate_model_paths, so no conversion step is needed. Megatron-FSDP + # follows the same model-config route because its weights are loaded + # directly from the HF snapshot rather than a torch_dist conversion. + # Refit-fed policies also derive architecture from HF because they never + # load the source checkpoint's serialized provider. from transformers import AutoConfig hf_config_overrides = config.get("hf_config_overrides", {}) or {} @@ -1159,6 +1204,7 @@ def setup_model_config( optimizer_path, load_main_params_from_ckpt, ckpt_cfg=None if skip_weight_load else config["megatron_cfg"].get("checkpoint"), + use_megatron_fsdp=use_megatron_fsdp, ) # Validate training configuration @@ -1296,13 +1342,15 @@ def _apply_parallelism_config(model_cfg: Any, config: PolicyConfig) -> None: model_cfg.context_parallel_size = config["megatron_cfg"]["context_parallel_size"] if model_cfg.context_parallel_size > 1: - # Either NeMo-RL does the packing+CP-sharding itself (classic mcore - # GPTModel path) OR the model does it internally (mbridge VLM wrappers - # like Qwen3VL, auto-detected at model build). Both paths require - # cu_seqlens to flow via PackedSeqParams, so sequence_packing must be on. - assert config["sequence_packing"]["enabled"], ( - "Sequence Packing must be enabled to use Context Parallelism with MCore." - ) + # Three CP layouts are supported: + # * sequence packing on, NeMo-RL packs + CP-shards (classic mcore + # GPTModel path), + # * sequence packing on, the model packs + CP-shards internally + # (mbridge VLM wrappers like Qwen3VL, auto-detected at model build), + # * sequence packing off, NeMo-RL CP-shards the [B, S] batch directly + # (process_microbatch's unpacked branch) -- used by cross-tokenizer + # distillation, whose per-sample teacher/alignment payload has no + # packed representation. assert not config["megatron_cfg"].get("use_fused_linear_logprobs", False), ( "Context Parallelism is not supported with linear CE fusion loss, please set use_fused_linear_logprobs to false" ) @@ -1974,6 +2022,7 @@ def _create_checkpoint_config( optimizer_path: Optional[str], load_main_params_from_ckpt: bool = False, ckpt_cfg: Optional[dict[str, Any]] = None, + use_megatron_fsdp: bool = False, ) -> CheckpointConfig: """Create checkpoint configurations. @@ -2002,6 +2051,10 @@ def _create_checkpoint_config( load_rng=False, load_main_params_from_ckpt=load_main_params_from_ckpt, ) + if use_megatron_fsdp: + # Megatron-FSDP checkpoints use DTensor/DCP and are not interchangeable + # with the ordinary Megatron torch_dist format. + kwargs["ckpt_format"] = "fsdp_dtensor" # Forward checkpoint knobs only when explicitly set in YAML; otherwise Megatron # Bridge's own CheckpointConfig defaults apply (the exemplar configs own the # values). async_save is presence-checked exactly like the sibling Bridge knobs @@ -2126,7 +2179,9 @@ def _create_megatron_config( "overlap_param_gather=false." ) + use_megatron_fsdp = _use_megatron_fsdp(config) dist_cfg = DistributedInitConfig() + dist_cfg.use_megatron_fsdp = use_megatron_fsdp # Bridge creates the model-parallel groups (TP/PP/DP/CP/EP) with an explicit # timeout from this field, so it must match the default group's timeout. timeout_minutes = _distributed_timeout_minutes(config) @@ -2137,6 +2192,13 @@ def _create_megatron_config( "use_gloo_process_groups" ] + ddp_config = config["megatron_cfg"]["distributed_data_parallel_config"] + optional_ddp_kwargs: dict[str, Any] = {} + if "reduce_scatter_with_fp32_accumulation" in ddp_config: + optional_ddp_kwargs["reduce_scatter_with_fp32_accumulation"] = ddp_config[ + "reduce_scatter_with_fp32_accumulation" + ] + return ConfigContainer( model=model_cfg, checkpoint=checkpoint_config, @@ -2160,6 +2222,7 @@ def _create_megatron_config( # we need to set average_in_collective=False with calculate_per_token_loss=T # otherwise, mcore throws an assertion error. average_in_collective=False, # Required with calculate_per_token_loss=True + use_megatron_fsdp=use_megatron_fsdp, use_distributed_optimizer=config["megatron_cfg"]["optimizer"][ "use_distributed_optimizer" ], @@ -2168,12 +2231,17 @@ def _create_megatron_config( ]["data_parallel_sharding_strategy"], reuse_grad_buf_for_mxfp8_param_ag=reuse_grad_buf_for_mxfp8_param_ag, fp8_param_gather=fp8_param_enabled, + **optional_ddp_kwargs, ), scheduler=SchedulerConfig(**config["megatron_cfg"]["scheduler"]), dataset=None, tokenizer=TokenizerConfig( tokenizer_type="HuggingFaceTokenizer", - tokenizer_model=hf_model_name, + # The model source and tokenizer source are intentionally independent. + # In particular, qualification jobs may load a converted Megatron + # checkpoint by its stable model id while pinning tokenization to an + # exact local HuggingFace snapshot. + tokenizer_model=config["tokenizer"]["name"], ), ) @@ -2238,6 +2306,42 @@ def draft_pre_wrap_hook(model: list[MegatronModule]) -> list[MegatronModule]: return draft_pre_wrap_hook +def _create_freeze_parameter_patterns_pre_wrap_hook( + freeze_patterns: list[str], +) -> Callable: + """Create a hook that freezes parameters whose names contain a pattern. + + This deliberately matches the AutoModel backend's substring semantics + (``pattern in parameter_name``). The hook runs after the built-in and + caller-supplied model-construction hooks so parameters added by PEFT, + draft, or value-head hooks are included before the optimizer is built. + """ + patterns = tuple(freeze_patterns) + + def freeze_parameter_patterns(megatron_model): + model_chunks = ( + megatron_model if isinstance(megatron_model, list) else [megatron_model] + ) + n_matched = 0 + n_frozen = 0 + for model_chunk in model_chunks: + for name, param in model_chunk.named_parameters(): + if any(pattern in name for pattern in patterns): + n_matched += 1 + if param.requires_grad: + param.requires_grad_(False) + n_frozen += 1 + + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + print( + f"Freeze patterns matched {n_matched} parameters; newly froze " + f"{n_frozen}: freeze_parameter_patterns={list(patterns)}" + ) + return megatron_model + + return freeze_parameter_patterns + + _BRIDGE_SIGNAL_HANDLER_PATCHED = False @@ -2450,7 +2554,10 @@ def setup_model_and_optimizer( ) pretrained_checkpoint_exists = ( megatron_cfg.checkpoint.pretrained_checkpoint is not None - and checkpoint_exists(megatron_cfg.checkpoint.pretrained_checkpoint) + and ( + checkpoint_exists(megatron_cfg.checkpoint.pretrained_checkpoint) + or is_hf_checkpoint_dir(megatron_cfg.checkpoint.pretrained_checkpoint) + ) ) preload_policy_from_pretrained_for_draft = ( draft_enabled @@ -2463,9 +2570,10 @@ def setup_model_and_optimizer( if policy_cfg["megatron_cfg"]["freeze_moe_router"]: def freeze_moe_router(megatron_model): - if not isinstance(megatron_model, list): - megatron_model = [megatron_model] - for model_module in megatron_model: + model_chunks = ( + megatron_model if isinstance(megatron_model, list) else [megatron_model] + ) + for model_module in model_chunks: # Handle both wrapped (Float16Module) and unwrapped models if isinstance(model_module, Float16Module): model_module = model_module.module @@ -2485,6 +2593,7 @@ def freeze_moe_router(megatron_model): for layer in model_module.decoder.layers: if hasattr(layer, "mlp") and hasattr(layer.mlp, "router"): layer.mlp.router.weight.requires_grad = False + return megatron_model mixed_precision_wrapper = MoEFloat16Module pre_wrap_hook.extend([freeze_moe_router]) @@ -2624,6 +2733,12 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]: if additional_pre_wrap_hooks: pre_wrap_hook.extend(additional_pre_wrap_hooks) + freeze_patterns = policy_cfg.get("freeze_parameter_patterns") or [] + if freeze_patterns: + pre_wrap_hook.append( + _create_freeze_parameter_patterns_pre_wrap_hook(freeze_patterns) + ) + # Model, optimizer, and learning rate. pg_collection = ProcessGroupCollection.use_mpu_process_groups() setattr(megatron_cfg.model, "_pg_collection", pg_collection) @@ -2634,6 +2749,7 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]: model = get_model( megatron_cfg.model, megatron_cfg.ddp, + use_megatron_fsdp=megatron_cfg.dist.use_megatron_fsdp, use_torch_fsdp2=megatron_cfg.dist.use_torch_fsdp2, overlap_param_gather_with_optimizer_step=megatron_cfg.optimizer.overlap_param_gather_with_optimizer_step, data_parallel_random_init=megatron_cfg.rng.data_parallel_random_init, @@ -2789,6 +2905,12 @@ def handle_model_import( # validate_model_paths() already confirmed both exist, so nothing to do. return + # Megatron-FSDP is initialized directly from the local HuggingFace + # snapshot by Bridge's checkpoint loader. There is no intermediate + # torch_dist conversion to create (or force-recreate). + if _use_megatron_fsdp(config): + return + force_reconvert = config["megatron_cfg"].get("force_reconvert_from_hf", False) if pt_checkpoint_exists and not force_reconvert: @@ -3143,7 +3265,7 @@ def finalize_megatron_setup( tokenizer_config = TokenizerConfig( tokenizer_type="HuggingFaceTokenizer", - tokenizer_model=hf_model_name, + tokenizer_model=config["tokenizer"]["name"], hf_tokenizer_kwargs={ "trust_remote_code": True, "use_fast": True, diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py index 45d514f2772..7387547ccdd 100644 --- a/nemo_rl/models/megatron/train.py +++ b/nemo_rl/models/megatron/train.py @@ -14,6 +14,7 @@ from collections import defaultdict from contextlib import contextmanager, nullcontext +from dataclasses import dataclass from functools import partial from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union, cast @@ -50,6 +51,10 @@ prepare_packed_loss_input, wrap_loss_fn_with_input_preparation, ) +from nemo_rl.algorithms.loss.csd import CSDLossFn +from nemo_rl.algorithms.loss.csd_packing import ( + CSDSequencePackingLossWrapper, +) from nemo_rl.algorithms.loss.draft import DEFAULT_DRAFT_TOKEN_CHUNK_SIZE from nemo_rl.algorithms.loss.interfaces import LossFunction from nemo_rl.algorithms.loss.utils import _pack_input_ids @@ -57,8 +62,11 @@ from nemo_rl.data.megatron_sft_packed import is_direct_packed_row from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.model_utils import ( + _get_tokens_on_this_cp_rank, allgather_cp_sharded_tensor, + cp_load_balanced_to_contiguous, distributed_vocab_topk, + distributed_vocab_topk_logz_force, from_parallel_logits_to_logprobs, from_parallel_logits_to_logprobs_packed_sequences, ) @@ -82,8 +90,14 @@ "LogprobsPostProcessor", "TeacherFullPayloadPostProcessor", "TopkLogitsPostProcessor", + "SparseTeacherLogitsPostProcessor", + "FullLogitsPostProcessor", ] +# Teacher rows per sparse top-k/log Z tile when ``logprob_chunk_size`` is unset; +# bounds the fp32 ``[B, rows, V/TP]`` working set at long context. +SPARSE_TEACHER_ROW_CHUNK_SIZE = 128 + def _prepare_padding_mask_for_model( model: GPTModel, @@ -155,6 +169,25 @@ def suspend_activation_offload_for_forward_only( model_config.fine_grained_activation_offloading = original_value +@dataclass(frozen=True) +class StreamedFullLogitsMetadata: + """Tensor-free record for one packed full-logit microbatch. + + The corresponding valid-prefix values have already been copied into the + persistent IPC slab. Keeping only scalar geometry in Megatron's + ``forward_data_store`` prevents it from retaining either the model output or + a view that can obscure which allocation owns the exported payload. + """ + + batch_size: int + logical_seq_len: int + local_vocab_size: int + dtype: torch.dtype + device: torch.device + storage_token_offset: int + stored_seq_lengths: tuple[int, ...] + + def model_forward( model: GPTModel, data_dict: BatchedDataDict[Any], @@ -162,9 +195,9 @@ def model_forward( position_ids: Optional[torch.Tensor], attention_mask: Optional[torch.Tensor], packed_seq_params: Optional[PackedSeqParams] = None, + padding_mask: Optional[torch.Tensor] = None, defer_fp32_logits: Optional[bool] = False, mtp_loss_mask: Optional[torch.Tensor] = None, - padding_mask: Optional[torch.Tensor] = None, straggler_timer: Optional[StragglerDetector] = None, use_fused_linear_logprobs: bool = False, labels_cp_sharded: Optional[torch.Tensor] = None, @@ -301,6 +334,8 @@ def forward_with_post_processing_fn( defer_fp32_logits: Optional[bool] = False, global_valid_seqs: Optional[torch.Tensor] = None, global_valid_toks: Optional[torch.Tensor] = None, + global_valid_kd_toks: Optional[torch.Tensor] = None, + global_valid_chunks_by_idx: Optional[dict[int, torch.Tensor]] = None, sampling_params: Optional[TrainingSamplingParams] = None, straggler_timer: Optional[StragglerDetector] = None, draft_model: Optional[MegatronModule] = None, @@ -324,6 +359,7 @@ def forward_with_post_processing_fn( defer_fp32_logits: Whether to defer FP32 conversion of logits global_valid_seqs: Global valid sequence count for loss normalization global_valid_toks: Global valid token count for loss normalization + global_valid_chunks_by_idx: Global valid alignment-chunk counts by teacher sampling_params: Sampling parameters (top-k, top-p, temperature) straggler_timer: Straggler detector for profiling the forward pass draft_model: Draft model for online draft model training @@ -348,10 +384,10 @@ def forward_with_post_processing_fn( position_ids = processed_mb.position_ids packed_seq_params = processed_mb.packed_seq_params cu_seqlens_padded = processed_mb.cu_seqlens_padded - mtp_loss_mask = processed_mb.mtp_loss_mask padding_mask = processed_mb.padding_mask labels_cp_sharded = processed_mb.labels_cp_sharded loss_mask_cp_sharded = processed_mb.loss_mask_cp_sharded + mtp_loss_mask = processed_mb.mtp_loss_mask routed_experts_cp_sharded = processed_mb.routed_experts_cp_sharded original_seq_length = processed_mb.original_seq_length media_token_validity_mask = processed_mb.media_token_validity_mask @@ -384,9 +420,9 @@ def forward_with_post_processing_fn( position_ids=position_ids, attention_mask=attention_mask, packed_seq_params=packed_seq_params, + padding_mask=padding_mask, defer_fp32_logits=defer_fp32_logits, mtp_loss_mask=mtp_loss_mask, - padding_mask=padding_mask, straggler_timer=straggler_timer, use_fused_linear_logprobs=use_fused_linear_logprobs, labels_cp_sharded=labels_cp_sharded, @@ -467,6 +503,8 @@ def forward_with_post_processing_fn( global_valid_seqs=global_valid_seqs, global_valid_toks=global_valid_toks, prepacked_loss_mask=loss_mask_cp_sharded, + global_valid_kd_toks=global_valid_kd_toks, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, ) elif isinstance(post_processing_fn, LogprobsPostProcessor): assert original_seq_length is not None @@ -496,6 +534,17 @@ def forward_with_post_processing_fn( cu_seqlens_padded=cu_seqlens_padded, original_seq_length=original_seq_length, ) + elif isinstance(post_processing_fn, SparseTeacherLogitsPostProcessor): + post_processing_fn_wrapped = post_processing_fn( + data_dict=data_dict, + cu_seqlens_padded=cu_seqlens_padded, + ) + elif isinstance(post_processing_fn, FullLogitsPostProcessor): + post_processing_fn_wrapped = post_processing_fn( + data_dict=data_dict, + cu_seqlens_padded=cu_seqlens_padded, + packed_seq_params=packed_seq_params, + ) else: raise TypeError( f"Unknown post-processing function type: {type(post_processing_fn)}" @@ -515,6 +564,8 @@ def megatron_forward_backward( defer_fp32_logits: Optional[bool] = False, global_valid_seqs: Optional[torch.Tensor] = None, global_valid_toks: Optional[torch.Tensor] = None, + global_valid_kd_toks: Optional[torch.Tensor] = None, + global_valid_chunks_by_idx: Optional[dict[int, torch.Tensor]] = None, sampling_params: Optional[TrainingSamplingParams] = None, straggler_timer: Optional[StragglerDetector] = None, draft_model: Optional[MegatronModule] = None, @@ -542,6 +593,7 @@ def megatron_forward_backward( defer_fp32_logits: Whether to skip the conversion of logits to fp32 global_valid_seqs: Global valid sequence count for loss normalization global_valid_toks: Global valid token count for loss normalization + global_valid_chunks_by_idx: Global valid alignment-chunk counts by teacher sampling_params: Sampling parameters (top-k, top-p, temperature) straggler_timer: Straggler detector for profiling the forward pass draft_model: Draft model for online draft model training @@ -559,6 +611,8 @@ def megatron_forward_backward( defer_fp32_logits=defer_fp32_logits, global_valid_seqs=global_valid_seqs, global_valid_toks=global_valid_toks, + global_valid_kd_toks=global_valid_kd_toks, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, sampling_params=sampling_params, straggler_timer=straggler_timer, draft_model=draft_model, @@ -648,6 +702,8 @@ def __call__( global_valid_seqs: Optional[torch.Tensor] = None, global_valid_toks: Optional[torch.Tensor] = None, prepacked_loss_mask: Optional[torch.Tensor] = None, + global_valid_kd_toks: Optional[torch.Tensor] = None, + global_valid_chunks_by_idx: Optional[dict[int, torch.Tensor]] = None, ) -> Callable[[torch.Tensor], Tuple[torch.Tensor, Dict[str, Any]]]: """Create a loss post-processing function for training. @@ -662,6 +718,7 @@ def __call__( global_valid_toks: Global valid token count for loss normalization prepacked_loss_mask: CP-local target-aligned mask when MCore returned per-token losses for a direct packed row + global_valid_chunks_by_idx: Global valid alignment-chunk counts by teacher Returns: Callable: Function that takes output tensor and returns (loss, metrics) tuple @@ -735,7 +792,10 @@ def _counteract_direct_mcore_loss_averaging( pack_sequences = self.cfg["sequence_packing"]["enabled"] if pack_sequences and packed_seq_params is not None: fuse_loss = self.cfg.get("sequence_packing", {}).get("fuse_loss", False) - if fuse_loss: + if isinstance(self.loss_fn, CSDLossFn): + wrapper_cls = CSDSequencePackingLossWrapper + prepare_fn = prepare_loss_input_wrapped + elif fuse_loss: # The fused path prepares loss via prepare_packed_loss_input and # cannot honor a custom prepare_fn (e.g. the value model's); guard # before reading packed metadata so misconfiguration fails clearly. @@ -826,12 +886,16 @@ def _counteract_direct_mcore_loss_averaging( defer_normalization=self.defer_draft_normalization, ) - loss_fn_wrapped = partial( - loss_fn_wrapped, - data=data_dict, - global_valid_seqs=global_valid_seqs, - global_valid_toks=global_valid_toks, - ) + loss_kwargs: dict[str, Any] = { + "data": data_dict, + "global_valid_seqs": global_valid_seqs, + "global_valid_toks": global_valid_toks, + } + if isinstance(self.loss_fn, CSDLossFn): + loss_kwargs["global_valid_kd_toks"] = global_valid_kd_toks + if isinstance(self.loss_fn, CSDLossFn) and global_valid_chunks_by_idx: + loss_kwargs["global_valid_chunks_by_idx"] = global_valid_chunks_by_idx + loss_fn_wrapped = partial(loss_fn_wrapped, **loss_kwargs) if self.cp_normalize: cp_size = get_context_parallel_world_size() @@ -858,6 +922,16 @@ def _counteract_mcore_loss_averaging(*args, **kwargs): return loss_fn_wrapped +def loss_returns_context_parallel_partial(loss_fn: LossFunction) -> bool: + """Whether ``loss_fn`` returns a CP-partial loss that must not be divided by CP. + + Owner-local CSD evaluation returns each context-parallel rank's share of the + loss; Megatron sums gradients over DP x CP, so the shares already add up to + the full loss. + """ + return isinstance(loss_fn, CSDLossFn) and loss_fn.owner_local + + def should_reduce_loss_across_context_parallel( data: BatchedDataDict[Any] | dict[str, Any], ) -> bool: @@ -1302,6 +1376,412 @@ def processor_fn_inner(output_tensor): return processor_fn_inner +class SparseTeacherLogitsPostProcessor: + """Build sparse teacher rows while keeping the CP sequence sharding. + + Unlike :class:`TopkLogitsPostProcessor`, this transport-oriented processor + never gathers the sequence axis and never copies results through CPU/Ray. + It computes the global top-k and exact log Z across TP only, excludes + padded LM-head columns, and exports a bounded sidecar holding every + realized label the loss requests from each row. The natural top-k stays + unchanged so the consumer can build an exact K-wide support per request. + + The teacher data must carry ``force_include_token_ids`` ``[B, T, F]`` + (``-1`` for unused slots) at the teacher tokenization width. + """ + + def __init__( + self, + cfg: PolicyConfig, + *, + k: int, + temperature: float, + real_vocab_size: int, + noise_filter_k: int, + ): + self.cfg = cfg + self.k = int(k) + self.temperature = float(temperature) + self.real_vocab_size = int(real_vocab_size) + self.noise_filter_k = int(noise_filter_k) + + def __call__( + self, + data_dict: BatchedDataDict[Any], + cu_seqlens_padded: Optional[torch.Tensor], + ) -> Callable[[torch.Tensor], Tuple[torch.Tensor, Dict[str, torch.Tensor]]]: + """Create a post-processing function that emits this rank's sparse rows. + + Args: + data_dict: Batched data dictionary with ``force_include_token_ids``. + cu_seqlens_padded: Unused; sequence packing is not supported. + + Returns: + Callable that maps the CP-local model output ``[B, S_local, V/TP]`` + to ``(dummy_loss, sparse_fields)`` with ``sparse_topk_logits`` / + ``sparse_topk_indices`` ``[B, S_local, k]``, ``sparse_log_z`` + ``[B, S_local]``, ``sparse_natural_tail_indices`` ``[B, S_local, 1]``, + and ``sparse_forced_logits`` / ``sparse_forced_indices`` / + ``sparse_forced_in_topk`` ``[B, S_local, F]``. + """ + del cu_seqlens_padded + if self.cfg["sequence_packing"]["enabled"]: + raise NotImplementedError( + "SparseTeacherLogitsPostProcessor does not support sequence packing." + ) + if "force_include_token_ids" not in data_dict: + raise KeyError( + "Sparse teacher IPC requires force_include_token_ids in teacher data." + ) + + force_ids_full = data_dict["force_include_token_ids"] + cp_size = get_context_parallel_world_size() + cp_rank = ( + torch.distributed.get_rank(get_context_parallel_group()) + if cp_size > 1 + else 0 + ) + # force_include_token_ids rides along at the teacher tokenization width + # and is exempt from the sequence-dim validation. Nothing else checks + # it, and _get_tokens_on_this_cp_rank floor-divides without asserting, + # so a width mismatch would silently misplace rows at CP > 1. + if cp_size > 1 and force_ids_full.shape[1] % (2 * cp_size) != 0: + raise ValueError( + "Sparse teacher force_include_token_ids must be padded to a " + f"multiple of 2*cp_size={2 * cp_size} so its load-balanced rows " + "match the teacher logits; got sequence length " + f"{force_ids_full.shape[1]}." + ) + force_ids_local = _get_tokens_on_this_cp_rank( + force_ids_full, cp_rank, cp_size, seq_dim=1 + ) + chunk_size = self.cfg.get("logprob_chunk_size") + if chunk_size is None: + chunk_size = SPARSE_TEACHER_ROW_CHUNK_SIZE + + def processor_fn_inner(output_tensor): + # The forced-label rows must line up 1:1 with this rank's logit + # rows; otherwise forced labels resolve against the wrong positions. + if force_ids_local.shape[1] != output_tensor.shape[1]: + raise ValueError( + "Sparse teacher forced-label rows do not match the CP-local " + f"logit rows: force={force_ids_local.shape[1]}, " + f"logits={output_tensor.shape[1]}, cp_size={cp_size}. Set " + "the teacher's make_sequence_length_divisible_by to a " + "multiple of its Megatron sequence padding factor." + ) + tp_grp = get_tensor_model_parallel_group() + tp_rank = get_tensor_model_parallel_rank() + vocab_shard_size = int(output_tensor.shape[-1]) + ( + values, + indices, + log_z, + natural_tail_indices, + forced_values, + forced_indices, + forced_in_topk, + ) = distributed_vocab_topk_logz_force( + output_tensor, + force_ids_local, + self.k, + tp_grp, + vocab_start_index=tp_rank * vocab_shard_size, + vocab_end_index=(tp_rank + 1) * vocab_shard_size, + real_vocab_size=self.real_vocab_size, + temperature=self.temperature, + noise_filter_k=self.noise_filter_k, + chunk_size=chunk_size, + ) + return output_tensor.new_zeros(()), { + "sparse_topk_logits": values, + "sparse_topk_indices": indices, + "sparse_log_z": log_z, + "sparse_natural_tail_indices": natural_tail_indices, + "sparse_forced_logits": forced_values, + "sparse_forced_indices": forced_indices, + "sparse_forced_in_topk": forced_in_topk, + } + + return processor_fn_inner + + +class FullLogitsPostProcessor: + """Export this rank's raw teacher logits (full local vocab shard, no reduction). + + Cross-tokenizer distillation ships the teacher's full-vocab logits to the + student and does all vocab reduction student-side, so each tensor-parallel + rank emits its local vocab shard untouched (fp32). Under context parallelism + each rank additionally re-emits its *contiguous* sequence slice (see below). + The IPC consumer reassembles the global ``[B, T_t, V_t]`` from the per-rank + shards. This is the Megatron counterpart of the Automodel + ``FullLogitsPostProcessor``. + + Packed THD outputs normally restore one dense logical row per input sample. + With ``packed_output_storage``, their valid prefixes are instead streamed + directly into a flat IPC slab. Under CP, each physical sequence is first + reconstructed from its native per-sequence head/tail shards, then this rank + emits the same contiguous logical sequence window used by the unpacked IPC + contract. + """ + + def __init__( + self, + cfg: PolicyConfig, + *, + packed_output_storage: Optional[torch.Tensor] = None, + ) -> None: + self.cfg = cfg + self.packed_output_storage = packed_output_storage + self._packed_output_token_cursor = 0 + if packed_output_storage is not None: + if not cfg["sequence_packing"]["enabled"]: + raise ValueError( + "Packed full-logit streaming storage requires sequence packing." + ) + if packed_output_storage.ndim != 2: + raise ValueError( + "Packed full-logit streaming storage must be two-dimensional, " + f"got shape={tuple(packed_output_storage.shape)}." + ) + if packed_output_storage.dtype != torch.float32: + raise TypeError( + "Packed full-logit streaming storage must use torch.float32, " + f"got {packed_output_storage.dtype}." + ) + + @property + def packed_output_token_cursor(self) -> int: + """Number of valid-prefix tokens committed to streaming storage.""" + return self._packed_output_token_cursor + + def __call__( + self, + data_dict: BatchedDataDict[Any], + cu_seqlens_padded: Optional[torch.Tensor], + packed_seq_params: Optional[PackedSeqParams] = None, + ) -> Callable[ + [torch.Tensor], + Tuple[ + torch.Tensor, + Dict[str, Union[torch.Tensor, StreamedFullLogitsMetadata]], + ], + ]: + """Create a post-processing function that emits raw full-vocab logit shards. + + Args: + data_dict: Batched data dictionary. + cu_seqlens_padded: Physical cumulative sequence boundaries. + packed_seq_params: Packed metadata carrying distinct raw and padded + cumulative boundaries. + + Returns: + Callable that maps the model output tensor to + ``(dummy_loss, {"full_logits": local_vocab_shard})`` where the shard + is ``[B, S_local_contiguous, V_local]`` in fp32. When + ``packed_output_storage`` was supplied at construction, packed values + are copied directly into that flat valid-prefix slab and + ``full_logits`` is tensor-free :class:`StreamedFullLogitsMetadata`. + """ + pack = self.cfg["sequence_packing"]["enabled"] + if pack and (packed_seq_params is None or cu_seqlens_padded is None): + raise ValueError( + "Packed full-logit export requires PackedSeqParams and padded " + "cumulative sequence boundaries." + ) + + def processor_fn_inner(output_tensor): + # fp32 for the consumer's precision-sensitive log_softmax / projection + # (KL math) and a dtype-consistent IPC buffer producer<->consumer. No + # vocab reduction: the local TP shard is emitted as-is and the student + # reassembles the full vocab. + logits = output_tensor.to(torch.float32) + if pack: + assert packed_seq_params is not None + assert cu_seqlens_padded is not None + raw_cu_seqlens = packed_seq_params.cu_seqlens_q + if raw_cu_seqlens is None: + raise ValueError( + "PackedSeqParams.cu_seqlens_q must preserve raw boundaries." + ) + batch_size, logical_seq_len = data_dict["input_ids"].shape[:2] + if ( + int(raw_cu_seqlens.numel()) != batch_size + 1 + or int(cu_seqlens_padded.numel()) != batch_size + 1 + ): + raise ValueError( + "Packed full-logit boundary cardinality does not match " + f"the logical batch size {batch_size}." + ) + cp_size = get_context_parallel_world_size() + if logical_seq_len % cp_size != 0: + raise ValueError( + f"Logical sequence width {logical_seq_len} must be " + f"divisible by context-parallel size {cp_size}." + ) + cp_rank = ( + torch.distributed.get_rank(get_context_parallel_group()) + if cp_size > 1 + else 0 + ) + local_logical_len = logical_seq_len // cp_size + logical_window_start = cp_rank * local_logical_len + logical_window_end = logical_window_start + local_logical_len + cp_padding_multiple = 2 * cp_size if cp_size > 1 else 1 + input_lengths = data_dict["input_lengths"] + cp_group = get_context_parallel_group() if cp_size > 1 else None + sample_slices: list[tuple[int, int, int, int]] = [] + stored_seq_lengths: list[int] = [] + for sample_index in range(batch_size): + raw_length = int( + ( + raw_cu_seqlens[sample_index + 1] + - raw_cu_seqlens[sample_index] + ).item() + ) + input_length = int(input_lengths[sample_index].item()) + if raw_length != input_length: + raise ValueError( + "Packed raw boundary disagrees with input_lengths " + f"for sample {sample_index}: raw={raw_length}, " + f"input_length={input_length}." + ) + padded_start = int(cu_seqlens_padded[sample_index].item()) + padded_end = int(cu_seqlens_padded[sample_index + 1].item()) + padded_length = padded_end - padded_start + if ( + padded_start % cp_padding_multiple != 0 + or padded_length % cp_padding_multiple != 0 + ): + raise ValueError( + "Packed padded boundaries must be divisible by the " + "load-balanced CP chunk count: " + f"sample={sample_index}, start={padded_start}, " + f"length={padded_length}, " + f"multiple={cp_padding_multiple}." + ) + local_start = padded_start // cp_size + local_length = padded_length // cp_size + if local_start + local_length > logits.shape[1]: + raise ValueError( + "Packed logit slice exceeds local output width for " + f"sample {sample_index}." + ) + overlap_start = logical_window_start + overlap_end = min(raw_length, logical_window_end) + stored_seq_len = max(0, overlap_end - overlap_start) + sample_slices.append( + (local_start, local_length, overlap_start, stored_seq_len) + ) + stored_seq_lengths.append(stored_seq_len) + + output_storage = self.packed_output_storage + storage_token_offset = self._packed_output_token_cursor + stored_tokens = sum(stored_seq_lengths) + next_token_cursor = storage_token_offset + stored_tokens + if output_storage is None: + unpacked = logits.new_zeros( + (batch_size, local_logical_len, logits.shape[-1]) + ) + else: + if ( + output_storage.shape[1] != logits.shape[-1] + or output_storage.device != logits.device + ): + raise ValueError( + "Packed full-logit streaming storage does not match " + "the model output vocabulary/device: " + f"storage={tuple(output_storage.shape)}/" + f"{output_storage.device}, logits={tuple(logits.shape)}/" + f"{logits.device}." + ) + if next_token_cursor > output_storage.shape[0]: + raise ValueError( + "Packed full-logit streaming output exceeds the " + "preallocated IPC slab: " + f"required={next_token_cursor}, " + f"capacity={output_storage.shape[0]}." + ) + + destination_token_cursor = storage_token_offset + for sample_index, ( + local_start, + local_length, + source_start, + stored_seq_len, + ) in enumerate(sample_slices): + physical = logits.narrow(1, local_start, local_length) + if cp_size > 1: + assert cp_group is not None + padded_length = local_length * cp_size + physical = allgather_cp_sharded_tensor( + physical, cp_group, seq_dim=1 + ) + if physical.shape[1] != padded_length: + physical = physical.reshape( + 1, padded_length, physical.shape[-1] + ) + if stored_seq_len > 0: + if output_storage is None: + destination_start = source_start - logical_window_start + unpacked[ + sample_index, + destination_start : destination_start + stored_seq_len, + ].copy_( + physical[ + 0, source_start : source_start + stored_seq_len + ] + ) + else: + next_destination = destination_token_cursor + stored_seq_len + output_storage[ + destination_token_cursor:next_destination + ].copy_( + physical[ + 0, source_start : source_start + stored_seq_len + ] + ) + destination_token_cursor = next_destination + del physical + + if output_storage is not None: + if destination_token_cursor != next_token_cursor: + raise RuntimeError( + "Packed full-logit streaming cursor did not consume " + "the validated microbatch geometry: " + f"cursor={destination_token_cursor}, " + f"expected={next_token_cursor}." + ) + self._packed_output_token_cursor = next_token_cursor + return output_tensor.new_zeros(()), { + "full_logits": StreamedFullLogitsMetadata( + batch_size=batch_size, + logical_seq_len=local_logical_len, + local_vocab_size=int(logits.shape[-1]), + dtype=logits.dtype, + device=logits.device, + storage_token_offset=storage_token_offset, + stored_seq_lengths=tuple(stored_seq_lengths), + ) + } + + return output_tensor.new_zeros(()), {"full_logits": unpacked} + + # mcore shards the seq dim load-balanced (each rank holds chunks + # ``cp_rank`` and ``2*cp-1-cp_rank`` of ``2*cp``), but the IPC + # consumer routes shards by a contiguous ``global_seq_start`` over + # the teacher CP group. Restore global order and emit this rank's + # contiguous slice, else a teacher_cp != student_cp pairing lands + # teacher data at the wrong seq positions in the consumer's dest. + if get_context_parallel_world_size() > 1: + logits = cp_load_balanced_to_contiguous( + logits, cp_group=get_context_parallel_group(), seq_dim=1 + ) + return output_tensor.new_zeros(()), {"full_logits": logits} + + return processor_fn_inner + + def aggregate_training_statistics( all_mb_metrics: List[Dict[str, Any]], losses: List[Union[float, torch.Tensor]], diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index e2dc97095fb..5caebfb8709 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -214,6 +214,17 @@ class AutomodelConfig(TypedDict): # Model config lora_cfg: NotRequired[LoRAConfig | LoRAConfigDisabled] automodel_kwargs: NotRequired[AutomodelKwargs] + # Dtype used when loading weights via from_pretrained. Valid values: + # "float32", "bfloat16", "float16". When omitted, optimizer-aware load + # dtype selection is preserved. Reference/teacher workers + # (init_optimizer=False) can set this to "bfloat16" / "float16" to halve + # init-time and resident parameter memory. Trainable workers MUST keep + # an explicit override at "float32" to preserve master-weight precision; non-float32 + # is rejected by validate_and_prepare_config when init_optimizer=True. + load_precision: NotRequired[str] + # Shard the model before loading its weights instead of loading the + # unwrapped model first. Lowers peak memory for large checkpoints. + shard_before_load: NotRequired[bool] # Runtime clear_cache_every_n_steps: NotRequired[int | None] checkpoint: NotRequired[AutomodelCheckpointConfig] @@ -224,6 +235,7 @@ class AutomodelConfig(TypedDict): class SequencePackingConfigDisabled(TypedDict): enabled: Literal[False] + fuse_loss: NotRequired[bool] class SequencePackingConfig(TypedDict): @@ -232,6 +244,7 @@ class SequencePackingConfig(TypedDict): # Not required because some algorithms like SFT don't calculate log probs logprob_mb_tokens: NotRequired[int] algorithm: str + sequence_length_round: NotRequired[int] # Preserve the packer's order (or omit for backward compatibility), or # execute each DP rank's assigned bins largest-first for allocator reuse. microbatch_order: NotRequired[Literal["packer", "largest_first"]] @@ -291,6 +304,13 @@ class MegatronOptimizerConfig(TypedDict): # distributed optimizer use_distributed_optimizer: bool use_precision_aware_optimizer: bool + # Precision-aware optimizer tensor dtypes. These are optional MCore + # pass-throughs; when omitted, MCore owns their defaults. + main_grads_dtype: NotRequired[str] + main_params_dtype: NotRequired[str] + exp_avg_dtype: NotRequired[str] + exp_avg_sq_dtype: NotRequired[str] + store_param_remainders: NotRequired[bool] clip_grad: float # knob to enable optimizer cpu offload optimizer_cpu_offload: bool @@ -316,8 +336,12 @@ class MegatronSchedulerConfig(TypedDict): class MegatronDDPConfig(TypedDict): grad_reduce_in_fp32: bool + # Communicate low-precision gradients while accumulating each reduced shard + # in FP32, then cast the result back to the gradient-buffer dtype. + reduce_scatter_with_fp32_accumulation: NotRequired[bool] overlap_grad_reduce: bool overlap_param_gather: bool + use_megatron_fsdp: NotRequired[bool] use_custom_fsdp: bool data_parallel_sharding_strategy: str @@ -720,6 +744,12 @@ class PolicyConfig(TypedDict): max_total_sequence_length: int # This sets the clipping norm for the AutomodelPolicyWorkers (Megatron's is called clip_grad) max_grad_norm: NotRequired[float | int | None] + # Supported by the Automodel and Megatron backends. + # Freeze (requires_grad=False) any parameter whose name contains one of these + # substrings. Use for submodules that never receive gradients (e.g. the + # nemotron_h multi-token-prediction head, ["mtp"], during distillation) to keep + # optimizer save/load symmetric and avoid checkpoint-resume key mismatches. + freeze_parameter_patterns: NotRequired[list[str] | None] refit_buffer_size_gb: NotRequired[float | int] # Run the policy offload lifecycle before non-colocated refit. offload_policy_before_refit: bool diff --git a/nemo_rl/models/policy/interfaces.py b/nemo_rl/models/policy/interfaces.py index 075efbd8cb8..817f54444f8 100644 --- a/nemo_rl/models/policy/interfaces.py +++ b/nemo_rl/models/policy/interfaces.py @@ -18,6 +18,7 @@ import torch from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.data.packing import LockstepPackingPlan from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.generation.interfaces import GenerationDatumSpec from nemo_rl.utils.timer import Timer @@ -129,6 +130,8 @@ def train( gbs: Optional[int] = None, mbs: Optional[int] = None, timer: Optional[Timer] = None, + packing_plan: Optional[LockstepPackingPlan] = None, + packing_side_id: Optional[str] = None, ) -> dict[str, Any]: """Train the policy on a global batch of data. @@ -138,6 +141,8 @@ def train( eval_mode: Whether to run in evaluation mode (no gradient updates) gbs: Global batch size override (if None, uses config default) mbs: Micro batch size override (if None, uses config default) + packing_plan: Optional controller-owned CSD lockstep packing plan. + packing_side_id: Side key in ``packing_plan``; provided together. """ pass diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index d7713fc0a66..585eb10109b 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -28,6 +28,7 @@ MEGATRON_SFT_PACKED_BATCH_FIELDS, is_direct_packed_row, ) +from nemo_rl.data.packing import LockstepPackingPlan from nemo_rl.distributed.batched_data_dict import ( BatchedDataDict, DynamicBatchingArgs, @@ -620,12 +621,33 @@ def _shard_for_logprob( unsorted_data_indices = None return sharded_data, unsorted_data_indices + def _validate_lockstep_plan_dp( + self, packing_plan: LockstepPackingPlan, packing_side_id: str + ) -> None: + """Fail before dispatch when a plan targets a different DP grid.""" + if packing_side_id not in packing_plan.sides: + raise ValueError( + f"Lockstep plan batch_uid={packing_plan.batch_uid} has no side " + f"{packing_side_id!r}; available sides are " + f"{tuple(packing_plan.sides)!r}." + ) + plan_dp_size = len(packing_plan.sides[packing_side_id].rank_bin_indices) + policy_dp_size = self.data_parallel_size + if plan_dp_size != policy_dp_size: + raise ValueError( + f"Lockstep plan batch_uid={packing_plan.batch_uid} side " + f"{packing_side_id!r} targets DP={plan_dp_size}, but this policy " + f"uses DP={policy_dp_size}." + ) + def _shard_for_train( self, data: BatchedDataDict[Any], batch_size: int, *, micro_batch_size: Optional[int] = None, + packing_plan: Optional[LockstepPackingPlan] = None, + packing_side_id: Optional[str] = None, ) -> list["SlicedDataDict"]: """Shard inputs for ``train``. @@ -636,7 +658,15 @@ def _shard_for_train( scalar metrics (no per-row outputs to reorder). """ dp_size = self.data_parallel_size + if (packing_plan is None) != (packing_side_id is None): + raise ValueError( + "packing_plan and packing_side_id must be provided together." + ) if is_direct_packed_row(data): + if packing_plan is not None: + raise ValueError( + "Direct packed rows cannot accept a lockstep packing plan." + ) missing = _DIRECT_PACKED_SFT_REQUIRED_KEYS.difference(data) if missing: raise ValueError( @@ -679,6 +709,26 @@ def _shard_for_train( ) for dp_rank in range(dp_size) ] + elif packing_plan is not None: + if not self.use_sequence_packing: + raise ValueError( + "A lockstep packing plan was supplied to a policy with " + "sequence packing disabled." + ) + if self.use_dynamic_batches: + raise ValueError( + "Lockstep sequence packing cannot be combined with dynamic batching." + ) + if batch_size != len(packing_plan.canonical_batch_item_ids): + raise ValueError( + f"train batch_size={batch_size} does not match lockstep plan " + f"logical size={len(packing_plan.canonical_batch_item_ids)}." + ) + assert packing_side_id is not None + self._validate_lockstep_plan_dp(packing_plan, packing_side_id) + sharded_data = data.shard_by_packing_plan( + packing_plan, side_id=packing_side_id + ) elif self.use_dynamic_batches: self.dynamic_batching_args["max_tokens_per_microbatch"] = self.cfg[ "dynamic_batching" @@ -873,6 +923,9 @@ def get_full_logits_ipc( data: BatchedDataDict[GenerationDatumSpec], micro_batch_size: Optional[int] = None, timer: Optional[Timer] = None, + *, + packing_plan: Optional[LockstepPackingPlan] = None, + packing_side_id: Optional[str] = None, ) -> list[dict[str, Any]]: """Ship the teacher's full-vocab logits to the student via CUDA IPC. @@ -891,19 +944,42 @@ def get_full_logits_ipc( :meth:`release_ipc_buffer` at the end of training / validation (and on error), not per call. - v0 limitation: no dynamic batching, no sequence packing. + Packed CSD callers must supply the controller-owned plan and side ID; + this method never recomputes membership locally. """ - if self.use_dynamic_batches or self.use_sequence_packing: + if self.use_dynamic_batches: raise NotImplementedError( - "get_full_logits_ipc does not support dynamic batching " - "or sequence packing in v0." + "get_full_logits_ipc does not support dynamic batching." + ) + if (packing_plan is None) != (packing_side_id is None): + raise ValueError( + "packing_plan and packing_side_id must be provided together." ) + if self.use_sequence_packing and packing_plan is None: + raise ValueError( + "Packed full-logits export requires the controller-owned " + "packing_plan and packing_side_id." + ) + if packing_plan is not None and not self.use_sequence_packing: + raise ValueError( + "A lockstep packing plan was supplied to a policy with sequence " + "packing disabled." + ) + if packing_plan is not None: + assert packing_side_id is not None + self._validate_lockstep_plan_dp(packing_plan, packing_side_id) dp_size = self.data_parallel_size with timer.time("get_full_logits_ipc/shard_data") if timer else nullcontext(): - sharded_data = data.shard_by_batch_size( # type: ignore - dp_size, - batch_size=None, - ) + if packing_plan is not None: + assert packing_side_id is not None + sharded_data = data.shard_by_packing_plan( + packing_plan, side_id=packing_side_id + ) + else: + sharded_data = data.shard_by_batch_size( # type: ignore + dp_size, + batch_size=None, + ) with timer.time("get_full_logits_ipc/submit") if timer else nullcontext(): futures = self.worker_group.run_all_workers_sharded_data( "get_full_logits_ipc", @@ -914,11 +990,134 @@ def get_full_logits_ipc( "tensor_parallel", "pipeline_parallel", ], - # Keep every TP × CP output; consumer routes via overlap. - output_is_replicated=["pipeline_parallel"], + # Keep every TP × CP × PP output; the consumer routes shards by + # their vocab/seq metadata. PP is deliberately NOT declared + # replicated: only the LAST pipeline stage holds logits, and + # "replicated" would return stage 0's result — which is empty. + # Non-last stages return no handles and drop out in + # aggregate_per_sample_handles. common_kwargs={"micro_batch_size": micro_batch_size}, ) worker_results = self.worker_group.get_all_worker_results(futures) + if packing_plan is not None: + return aggregate_per_sample_handles( + worker_results, + canonical_batch_item_ids=packing_plan.canonical_batch_item_ids, + ) + return aggregate_per_sample_handles(worker_results) + + def get_topk_logits_ipc( + self, + data: BatchedDataDict[GenerationDatumSpec], + *, + k: int, + temperature: float, + vocab_size: Optional[int] = None, + micro_batch_size: Optional[int] = None, + support_mode: str, + gt_filter_topk: Optional[int] = None, + timer: Optional[Timer] = None, + ) -> list[dict[str, Any]]: + """Ship sparse teacher logits, ids, and full-vocab logZ via CUDA IPC. + + Tensor-parallel workers reduce to the same global top-k support; context + parallel workers retain distinct contiguous sequence shards, which are + grouped per sample for reconstruction by the student loss worker. + """ + automodel_cfg = self.cfg["automodel_cfg"] + if not automodel_cfg["enabled"]: + raise NotImplementedError( + "get_topk_logits_ipc currently supports only Automodel workers." + ) + if self.use_dynamic_batches or self.use_sequence_packing: + raise NotImplementedError( + "get_topk_logits_ipc does not support dynamic batching or " + "sequence packing." + ) + with timer.time("get_topk_logits_ipc/shard_data") if timer else nullcontext(): + sharded_data = data.shard_by_batch_size( + self.data_parallel_size, + batch_size=None, + ) + with timer.time("get_topk_logits_ipc/submit") if timer else nullcontext(): + futures = self.worker_group.run_all_workers_sharded_data( + "get_topk_logits_ipc", + data=sharded_data, + in_sharded_axes=["data_parallel"], + replicate_on_axes=[ + "context_parallel", + "tensor_parallel", + "pipeline_parallel", + ], + output_is_replicated=["tensor_parallel", "pipeline_parallel"], + common_kwargs={ + "k": k, + "temperature": temperature, + "vocab_size": vocab_size, + "micro_batch_size": micro_batch_size, + "support_mode": support_mode, + "gt_filter_topk": gt_filter_topk, + }, + ) + worker_results = self.worker_group.get_all_worker_results(futures) + return aggregate_per_sample_handles(worker_results) + + def get_sparse_logits_ipc( + self, + data: BatchedDataDict[GenerationDatumSpec], + *, + k: int, + temperature: float, + real_vocab_size: int, + noise_filter_k: int, + micro_batch_size: Optional[int] = None, + timer: Optional[Timer] = None, + ) -> list[dict[str, Any]]: + """Ship sparse teacher rows of each CP shard via CUDA IPC (Megatron). + + TP ranks cooperate to form the exact top-k and log Z, but only TP rank + 0 exports each CP shard, so every sample keeps one sparse record per + teacher CP rank. ``data`` must carry ``force_include_token_ids`` + ``[B, T_t, F]``. Returns the global-batch-ordered list produced by + :func:`aggregate_per_sample_handles`. The worker storage persists across + calls and is freed by :meth:`release_ipc_buffer`. + """ + if "megatron_cfg" not in self.cfg or not self.cfg["megatron_cfg"]["enabled"]: + raise NotImplementedError( + "get_sparse_logits_ipc currently supports only Megatron workers." + ) + if self.use_dynamic_batches or self.use_sequence_packing: + raise NotImplementedError( + "get_sparse_logits_ipc does not support dynamic batching or " + "sequence packing." + ) + with timer.time("get_sparse_logits_ipc/shard_data") if timer else nullcontext(): + sharded_data = data.shard_by_batch_size( # type: ignore + self.data_parallel_size, + batch_size=None, + ) + with timer.time("get_sparse_logits_ipc/submit") if timer else nullcontext(): + futures = self.worker_group.run_all_workers_sharded_data( + "get_sparse_logits_ipc", + data=sharded_data, + in_sharded_axes=["data_parallel"], + replicate_on_axes=[ + "context_parallel", + "tensor_parallel", + "pipeline_parallel", + ], + # Keep every CP output; only TP rank 0 on the last pipeline + # stage returns records, and empty results drop out in + # aggregate_per_sample_handles. + common_kwargs={ + "k": int(k), + "temperature": float(temperature), + "real_vocab_size": int(real_vocab_size), + "noise_filter_k": int(noise_filter_k), + "micro_batch_size": micro_batch_size, + }, + ) + worker_results = self.worker_group.get_all_worker_results(futures) return aggregate_per_sample_handles(worker_results) def release_ipc_buffer(self) -> None: @@ -935,6 +1134,8 @@ def train( mbs: Optional[int] = None, timer: Optional[Timer] = None, check_dim_skip_keys: Optional[Iterable[str]] = None, + packing_plan: Optional[LockstepPackingPlan] = None, + packing_side_id: Optional[str] = None, ) -> dict[str, Any]: """Train the policy on a batch of data with a given loss function. @@ -944,19 +1145,21 @@ def train( pre-flight check. Used by cross-tokenizer distillation to pass through teacher / alignment auxiliaries that ride on the same data dict. + packing_plan: Controller-owned shared membership and side geometry. + packing_side_id: Side key within ``packing_plan``; supplied together + with ``packing_plan`` for CSD lockstep packing. """ batch_size = gbs or self.cfg["train_global_batch_size"] micro_batch_size = mbs or self.cfg["train_micro_batch_size"] # Shard and replicate the batch with timer.time("policy_training/sharding_data") if timer else nullcontext(): - if is_direct_packed_row(data): - sharded_data = self._shard_for_train( - data, - batch_size, - micro_batch_size=micro_batch_size, - ) - else: - sharded_data = self._shard_for_train(data, batch_size) + sharded_data = self._shard_for_train( + data, + batch_size, + micro_batch_size=micro_batch_size, + packing_plan=packing_plan, + packing_side_id=packing_side_id, + ) self._report_sharded_payload(sharded_data, "policy_train") if self.flops_tracker is not None: diff --git a/nemo_rl/models/policy/utils.py b/nemo_rl/models/policy/utils.py index 5fbebefd674..306f990c309 100644 --- a/nemo_rl/models/policy/utils.py +++ b/nemo_rl/models/policy/utils.py @@ -18,7 +18,7 @@ import warnings from datetime import timedelta from enum import Enum -from typing import TYPE_CHECKING, Any, Iterable, Optional, cast +from typing import TYPE_CHECKING, Any, Iterable, Mapping, Optional, cast import torch import torch.distributed as dist @@ -38,6 +38,8 @@ if TYPE_CHECKING: from nemo_rl.models.policy import PolicyConfig +DENSE_TEACHER_IPC_FLAT_LAYOUT = "flat_valid_prefix_v1" + # Environment for bitwise-reproducible policy workers. These are read when the # worker creates its cuBLAS handles and selects Transformer Engine kernels, so # they are set in the worker's runtime environment, before the process starts. @@ -363,14 +365,24 @@ def ensure_teacher_ipc_buffer( vocab_size: int, dtype: torch.dtype, device: torch.device, + capacity_seq_len: Optional[int] = None, ) -> tuple[torch.Tensor, tuple[Any, ...]]: """Lazy-alloc / grow ``[N_mb, B, T, V]`` teacher-logits IPC storage. + ``capacity_seq_len`` reserves a sequence high-water mark when the storage + is (re)allocated, so a later, slightly longer batch does not allocate a + replacement while the previous CUDA-IPC backing and the current forward + outputs are both live. The exported metadata still describes ``seq_len``; + the unused capacity is never read by the loss. + Returns the (possibly reallocated) ``(storage, handle)``. Reallocates and re-exports the IPC handle whenever any dim of the requested shape exceeds the current storage, or dtype/device changed; otherwise the existing storage and cached handle are returned unchanged. """ + allocation_seq_len = max( + seq_len, capacity_seq_len if capacity_seq_len is not None else seq_len + ) needs_realloc = ( storage is None or storage.shape[0] < num_microbatches @@ -382,7 +394,7 @@ def ensure_teacher_ipc_buffer( ) if needs_realloc: storage = torch.empty( - (num_microbatches, batch_size, seq_len, vocab_size), + (num_microbatches, batch_size, allocation_seq_len, vocab_size), dtype=dtype, device=device, ) @@ -391,26 +403,770 @@ def ensure_teacher_ipc_buffer( return storage, handle +def ensure_teacher_ipc_row_buffer( + storage: Optional[torch.Tensor], + total_rows: int, + seq_len: int, + vocab_size: int, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Lazy-allocate a compact ``[sum(B_mb), T, V]`` teacher IPC slab. + + Only the row capacity may be reused at a larger size. Sequence/vocabulary + geometry, dtype, and device must match exactly so every exported bin view is + contiguous and preserves the existing ``[1, B_mb, T, V]`` IPC contract. + """ + if total_rows <= 0 or seq_len <= 0 or vocab_size <= 0: + raise ValueError( + "Dense teacher IPC dimensions must be positive, got " + f"rows={total_rows}, seq_len={seq_len}, vocab_size={vocab_size}." + ) + needs_realloc = ( + storage is None + or storage.ndim != 3 + or storage.shape[0] < total_rows + or tuple(storage.shape[1:]) != (seq_len, vocab_size) + or storage.dtype != dtype + or storage.device != device + ) + if needs_realloc: + storage = torch.empty( + (total_rows, seq_len, vocab_size), + dtype=dtype, + device=device, + ) + assert storage is not None + return storage + + +def partition_teacher_ipc_row_buffer( + storage: torch.Tensor, batch_sizes: Iterable[int] +) -> list[tuple[int, torch.Tensor]]: + """Return compact row offsets and legacy-compatible 4-D bin views. + + Each view has shape ``[1, B_mb, T, V]`` and addresses a disjoint contiguous + range of the shared row slab. Exporting the views through CUDA IPC therefore + keeps existing consumers' ``buf_idx=0`` / bin-local sample indexing and + zero-copy fast path while avoiding ``N_mb * max(B_mb)`` padding rows. + """ + if storage.ndim != 3: + raise ValueError( + "Dense teacher IPC row storage must be three-dimensional, got " + f"shape={tuple(storage.shape)}." + ) + normalized_batch_sizes = [int(batch_size) for batch_size in batch_sizes] + if any(batch_size <= 0 for batch_size in normalized_batch_sizes): + raise ValueError( + "Dense teacher IPC microbatch cardinalities must be positive, got " + f"{normalized_batch_sizes}." + ) + total_rows = sum(normalized_batch_sizes) + if total_rows > storage.shape[0]: + raise ValueError( + "Dense teacher IPC row partitions exceed storage capacity: " + f"rows={total_rows}, capacity={storage.shape[0]}." + ) + + partitions: list[tuple[int, torch.Tensor]] = [] + row_offset = 0 + for batch_size in normalized_batch_sizes: + next_offset = row_offset + batch_size + partitions.append((row_offset, storage[row_offset:next_offset].unsqueeze(0))) + row_offset = next_offset + return partitions + + +def extract_teacher_ipc_valid_lengths( + data: Mapping[str, Any], + expected_batch_size: int, + *, + full_seq_len: int, + compact_padding: bool, +) -> list[int]: + """Return the logical valid prefix length for every teacher row. + + Packed :class:`FullLogitsPostProcessor` uses ``input_lengths`` as its + authoritative raw boundary. It either zero-fills the remainder of a logical + dense row or streams only the valid prefix; the compact consumer reconstructs + the omitted suffix as zero in both cases. The unpacked path does not promise + zero logits at padding positions, so it deliberately retains the complete + logical width. + """ + if expected_batch_size <= 0 or full_seq_len <= 0: + raise ValueError( + "Dense teacher IPC row geometry must be positive, got " + f"rows={expected_batch_size}, full_seq_len={full_seq_len}." + ) + if not compact_padding: + return [full_seq_len] * expected_batch_size + + values = data.get("input_lengths") + if values is None: + raise ValueError( + "Packed dense teacher IPC requires input_lengths on every logical " + "microbatch row." + ) + if isinstance(values, torch.Tensor): + if values.ndim != 1: + raise ValueError( + "input_lengths must be a one-dimensional tensor, got shape " + f"{tuple(values.shape)}." + ) + if torch.is_floating_point(values) or torch.is_complex(values): + raise TypeError( + "input_lengths must use an integer dtype for dense teacher IPC, " + f"got {values.dtype}." + ) + normalized = [int(value) for value in values.detach().cpu().tolist()] + elif isinstance(values, (list, tuple)): + normalized = [int(value) for value in values] + else: + raise TypeError( + "input_lengths must be a one-dimensional tensor, list, or tuple, " + f"got {type(values).__name__}." + ) + if len(normalized) != expected_batch_size: + raise ValueError( + "input_lengths cardinality does not match the logical microbatch: " + f"lengths={len(normalized)}, rows={expected_batch_size}." + ) + invalid = [value for value in normalized if value < 0 or value > full_seq_len] + if invalid: + raise ValueError( + "Dense teacher IPC input_lengths must lie within the logical " + f"sequence width [0, {full_seq_len}], got {invalid}." + ) + return normalized + + +def localize_teacher_ipc_valid_lengths( + valid_lengths: Iterable[int], + *, + full_seq_len: int, + cp_rank: int, + cp_size: int, +) -> list[int]: + """Intersect global valid prefixes with one teacher CP rectangle.""" + if cp_size <= 0 or cp_rank < 0 or cp_rank >= cp_size: + raise ValueError( + f"Invalid context-parallel coordinates: rank={cp_rank}, size={cp_size}." + ) + if full_seq_len <= 0 or full_seq_len % cp_size != 0: + raise ValueError( + f"Logical dense teacher width {full_seq_len} must be positive and " + f"divisible by context-parallel size {cp_size}." + ) + local_seq_len = full_seq_len // cp_size + global_seq_start = cp_rank * local_seq_len + global_seq_end = global_seq_start + local_seq_len + localized: list[int] = [] + for value in valid_lengths: + valid_seq_len = int(value) + if valid_seq_len < 0 or valid_seq_len > full_seq_len: + raise ValueError( + "Dense teacher IPC valid length must lie within the logical " + f"sequence width [0, {full_seq_len}], got {valid_seq_len}." + ) + localized.append(max(0, min(valid_seq_len, global_seq_end) - global_seq_start)) + return localized + + +def can_reuse_teacher_ipc_token_buffer( + storage: Optional[torch.Tensor], + *, + total_tokens: int, + vocab_size: int, + dtype: torch.dtype, + device: torch.device, +) -> bool: + """Return whether ``storage`` can hold one compact teacher IPC payload.""" + if total_tokens < 0 or vocab_size <= 0: + raise ValueError( + "Dense teacher IPC token geometry is invalid, got " + f"tokens={total_tokens}, vocab_size={vocab_size}." + ) + return ( + storage is not None + and storage.ndim == 2 + and storage.shape[0] >= total_tokens + and storage.shape[1] == vocab_size + and storage.dtype == dtype + and storage.device == device + ) + + +def ensure_teacher_ipc_token_buffer( + storage: Optional[torch.Tensor], + *, + total_tokens: int, + vocab_size: int, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Lazy-allocate a flat ``[sum(valid_local_T), V]`` teacher IPC slab.""" + if not can_reuse_teacher_ipc_token_buffer( + storage, + total_tokens=total_tokens, + vocab_size=vocab_size, + dtype=dtype, + device=device, + ): + storage = torch.empty( + (total_tokens, vocab_size), + dtype=dtype, + device=device, + ) + assert storage is not None + return storage + + +def partition_teacher_ipc_token_buffer( + storage: torch.Tensor, + valid_lengths_by_microbatch: Iterable[Iterable[int]], +) -> list[list[tuple[int, torch.Tensor]]]: + """Return token offsets and allocation-free row views into a flat slab.""" + if storage.ndim != 2: + raise ValueError( + "Dense teacher IPC token storage must be two-dimensional, got " + f"shape={tuple(storage.shape)}." + ) + normalized = [ + [int(valid_length) for valid_length in microbatch_lengths] + for microbatch_lengths in valid_lengths_by_microbatch + ] + if any(valid_length < 0 for lengths in normalized for valid_length in lengths): + raise ValueError( + "Dense teacher IPC local valid lengths must be non-negative, got " + f"{normalized}." + ) + total_tokens = sum(sum(lengths) for lengths in normalized) + if total_tokens > storage.shape[0]: + raise ValueError( + "Dense teacher IPC token partitions exceed storage capacity: " + f"tokens={total_tokens}, capacity={storage.shape[0]}." + ) + + partitions: list[list[tuple[int, torch.Tensor]]] = [] + token_offset = 0 + for microbatch_lengths in normalized: + microbatch_partitions: list[tuple[int, torch.Tensor]] = [] + for valid_length in microbatch_lengths: + next_offset = token_offset + valid_length + microbatch_partitions.append( + (token_offset, storage[token_offset:next_offset]) + ) + token_offset = next_offset + partitions.append(microbatch_partitions) + return partitions + + +def validate_compact_teacher_ipc_handle( + handle: Mapping[str, Any], +) -> Optional[tuple[int, int, int, int]]: + """Validate compact dense-IPC metadata and return its physical geometry. + + Legacy rectangular handles have no ``ipc_layout`` field and return + ``None``. Compact handles return ``(token_offset, stored_seq_len, + used_tokens, local_vocab_size)``. Unknown or partial schemas fail loudly + so a malformed handle cannot silently read another row's logits. + """ + layout = handle.get("ipc_layout") + if layout is None: + return None + if layout != DENSE_TEACHER_IPC_FLAT_LAYOUT: + raise ValueError(f"Unsupported dense teacher IPC layout {layout!r}.") + + required_fields = { + "actual_shape", + "dtype", + "full_seq_len", + "global_seq_start", + "payload_ipc", + "storage_capacity_tokens", + "storage_shape", + "storage_token_offset", + "storage_used_tokens", + "stored_seq_len", + "valid_seq_len", + } + missing_fields = required_fields - handle.keys() + if missing_fields: + raise ValueError( + "Compact dense teacher IPC handle is missing metadata " + f"{sorted(missing_fields)}." + ) + + actual_shape = handle["actual_shape"] + storage_shape = handle["storage_shape"] + if ( + not isinstance(actual_shape, (list, tuple, torch.Size)) + or len(actual_shape) != 2 + ): + raise ValueError( + "Compact dense teacher IPC actual_shape must be two-dimensional, " + f"got {actual_shape!r}." + ) + if ( + not isinstance(storage_shape, (list, tuple, torch.Size)) + or len(storage_shape) != 2 + ): + raise ValueError( + "Compact dense teacher IPC storage_shape must be two-dimensional, " + f"got {storage_shape!r}." + ) + + local_seq_len, local_vocab_size = (int(value) for value in actual_shape) + capacity_tokens, storage_vocab_size = (int(value) for value in storage_shape) + full_seq_len = int(handle["full_seq_len"]) + global_seq_start = int(handle["global_seq_start"]) + valid_seq_len = int(handle["valid_seq_len"]) + stored_seq_len = int(handle["stored_seq_len"]) + token_offset = int(handle["storage_token_offset"]) + used_tokens = int(handle["storage_used_tokens"]) + declared_capacity = int(handle["storage_capacity_tokens"]) + + if ( + local_seq_len <= 0 + or local_vocab_size <= 0 + or full_seq_len <= 0 + or global_seq_start < 0 + or global_seq_start + local_seq_len > full_seq_len + ): + raise ValueError( + "Compact dense teacher IPC has invalid logical sequence/vocabulary " + f"geometry: actual_shape={tuple(actual_shape)}, global_seq_start=" + f"{global_seq_start}, full_seq_len={full_seq_len}." + ) + if storage_vocab_size != local_vocab_size or capacity_tokens < 0: + raise ValueError( + "Compact dense teacher IPC storage shape disagrees with the logical " + f"vocabulary shard: storage_shape={tuple(storage_shape)}, " + f"actual_shape={tuple(actual_shape)}." + ) + if declared_capacity != capacity_tokens: + raise ValueError( + "Compact dense teacher IPC storage capacity metadata disagrees with " + f"storage_shape: declared={declared_capacity}, shape={capacity_tokens}." + ) + if valid_seq_len < 0 or valid_seq_len > full_seq_len: + raise ValueError( + "Compact dense teacher IPC valid_seq_len lies outside the logical " + f"sequence: valid={valid_seq_len}, full={full_seq_len}." + ) + expected_stored_seq_len = max( + 0, + min(valid_seq_len, global_seq_start + local_seq_len) - global_seq_start, + ) + if stored_seq_len != expected_stored_seq_len: + raise ValueError( + "Compact dense teacher IPC stored_seq_len disagrees with the valid " + f"CP overlap: stored={stored_seq_len}, expected=" + f"{expected_stored_seq_len}." + ) + if ( + used_tokens < 0 + or used_tokens > capacity_tokens + or token_offset < 0 + or token_offset > used_tokens + or token_offset + stored_seq_len > used_tokens + ): + raise ValueError( + "Compact dense teacher IPC token range is outside storage bounds: " + f"offset={token_offset}, stored={stored_seq_len}, used={used_tokens}, " + f"capacity={capacity_tokens}." + ) + if stored_seq_len > 0 and handle["payload_ipc"] is None: + raise ValueError( + "Compact dense teacher IPC omitted payload_ipc for a non-empty row." + ) + return token_offset, stored_seq_len, used_tokens, local_vocab_size + + +def extract_batch_item_ids( + data: Mapping[str, Any], + expected_batch_size: int, + *, + required: bool, +) -> list[int] | None: + """Read occurrence identities from one logical microbatch. + + Packed teacher IPC must be routed by occurrence identity rather than by a + bin-local position. The helper accepts the tensor or list representation + used by :class:`BatchedDataDict`, validates cardinality, and deliberately + does not synthesize identities when the controller did not provide them. + """ + values = data.get("batch_item_id") + if values is None: + if required: + raise ValueError( + "Packed dense teacher IPC requires batch_item_id on every " + "logical microbatch row." + ) + return None + if isinstance(values, torch.Tensor): + if values.ndim != 1: + raise ValueError( + "batch_item_id must be a one-dimensional tensor, got shape " + f"{tuple(values.shape)}." + ) + normalized = [int(value) for value in values.detach().cpu().tolist()] + else: + normalized = [int(value) for value in values] + if len(normalized) != expected_batch_size: + raise ValueError( + "batch_item_id cardinality does not match the logical microbatch: " + f"ids={len(normalized)}, rows={expected_batch_size}." + ) + if len(set(normalized)) != len(normalized): + raise ValueError( + f"batch_item_id values must be unique within a microbatch, got {normalized}." + ) + return normalized + + +def get_dense_ipc_sequence_layout( + data: Mapping[str, Any], *, cp_rank: int, cp_size: int +) -> tuple[int, int, int]: + """Return logical dense-IPC ``(full, local, start)`` sequence geometry. + + Packed backend schedulers operate on a physical bin width, which can be + much larger than the original per-sample token rectangle restored by the + full-logit postprocessor. IPC metadata must describe that logical rectangle + rather than the packed scheduler width so teacher masks and alignment axes + agree with reconstruction. + """ + if cp_size <= 0 or cp_rank < 0 or cp_rank >= cp_size: + raise ValueError( + f"Invalid context-parallel coordinates: rank={cp_rank}, size={cp_size}." + ) + input_ids = data.get("input_ids") + if not isinstance(input_ids, torch.Tensor) or input_ids.ndim < 2: + raise ValueError( + "Dense teacher IPC requires input_ids with a logical sequence axis." + ) + full_seq_len = int(input_ids.shape[1]) + if full_seq_len % cp_size != 0: + raise ValueError( + f"Logical dense teacher width {full_seq_len} must be divisible by " + f"context-parallel size {cp_size}." + ) + local_seq_len = full_seq_len // cp_size + return full_seq_len, local_seq_len, cp_rank * local_seq_len + + +def _validate_dense_ipc_shard_coverage( + batch_item_id: int, teacher_shards: list[dict[str, Any]] +) -> None: + """Validate that one sample has a complete, non-overlapping TP/CP grid.""" + required_fields = { + "actual_shape", + "cp_rank", + "cp_size", + "full_seq_len", + "full_vocab_size", + "global_seq_start", + "tp_rank", + "tp_size", + "vocab_end_index", + "vocab_start_index", + } + first = teacher_shards[0] + missing_fields = required_fields - first.keys() + if missing_fields: + raise ValueError( + f"Dense teacher IPC shard for batch_item_id={batch_item_id} is " + f"missing coverage metadata {sorted(missing_fields)}." + ) + + tp_size = int(first["tp_size"]) + cp_size = int(first["cp_size"]) + full_seq_len = int(first["full_seq_len"]) + full_vocab_size = int(first["full_vocab_size"]) + if tp_size <= 0 or cp_size <= 0: + raise ValueError( + f"Dense teacher IPC for batch_item_id={batch_item_id} has invalid " + f"parallel sizes tp_size={tp_size}, cp_size={cp_size}." + ) + if full_seq_len <= 0 or full_vocab_size <= 0: + raise ValueError( + f"Dense teacher IPC for batch_item_id={batch_item_id} has invalid " + f"logical shape ({full_seq_len}, {full_vocab_size})." + ) + + coordinates: set[tuple[int, int]] = set() + sequence_ranges: dict[int, tuple[int, int]] = {} + vocab_ranges: dict[int, tuple[int, int]] = {} + layouts: set[Optional[str]] = set() + valid_seq_lengths: set[int] = set() + compact_geometry_by_cp: dict[int, tuple[int, int, int, int]] = {} + for shard in teacher_shards: + missing_fields = required_fields - shard.keys() + if missing_fields: + raise ValueError( + f"Dense teacher IPC shard for batch_item_id={batch_item_id} is " + f"missing coverage metadata {sorted(missing_fields)}." + ) + shard_tp_size = int(shard["tp_size"]) + shard_cp_size = int(shard["cp_size"]) + shard_full_seq_len = int(shard["full_seq_len"]) + shard_full_vocab_size = int(shard["full_vocab_size"]) + if ( + shard_tp_size != tp_size + or shard_cp_size != cp_size + or shard_full_seq_len != full_seq_len + or shard_full_vocab_size != full_vocab_size + ): + raise ValueError( + "Dense teacher IPC shards disagree on topology or logical shape " + f"for batch_item_id={batch_item_id}." + ) + + tp_rank = int(shard["tp_rank"]) + cp_rank = int(shard["cp_rank"]) + if tp_rank < 0 or tp_rank >= tp_size or cp_rank < 0 or cp_rank >= cp_size: + raise ValueError( + f"Dense teacher IPC shard for batch_item_id={batch_item_id} has " + f"out-of-range coordinate (tp={tp_rank}, cp={cp_rank}) for " + f"topology ({tp_size}, {cp_size})." + ) + coordinate = (tp_rank, cp_rank) + if coordinate in coordinates: + raise ValueError( + f"Dense teacher IPC has duplicate shard coordinate {coordinate} " + f"for batch_item_id={batch_item_id}." + ) + coordinates.add(coordinate) + layouts.add(shard.get("ipc_layout")) + + actual_shape = shard["actual_shape"] + if ( + not isinstance(actual_shape, (list, tuple, torch.Size)) + or len(actual_shape) != 2 + ): + raise ValueError( + f"Dense teacher IPC shard for batch_item_id={batch_item_id} must " + f"have two-dimensional actual_shape, got {actual_shape!r}." + ) + local_seq_len, local_vocab_size = (int(value) for value in actual_shape) + seq_start = int(shard["global_seq_start"]) + seq_end = seq_start + local_seq_len + vocab_start = int(shard["vocab_start_index"]) + vocab_end = int(shard["vocab_end_index"]) + if ( + local_seq_len <= 0 + or local_vocab_size <= 0 + or seq_start < 0 + or seq_end > full_seq_len + or vocab_start < 0 + or vocab_end > full_vocab_size + or vocab_end - vocab_start != local_vocab_size + ): + raise ValueError( + f"Dense teacher IPC shard {coordinate} for " + f"batch_item_id={batch_item_id} has invalid rectangle " + f"seq=[{seq_start}, {seq_end}), vocab=[{vocab_start}, {vocab_end}), " + f"actual_shape={tuple(actual_shape)}, full_shape=" + f"({full_seq_len}, {full_vocab_size})." + ) + + sequence_range = (seq_start, seq_end) + prior_sequence_range = sequence_ranges.setdefault(cp_rank, sequence_range) + if prior_sequence_range != sequence_range: + raise ValueError( + f"Dense teacher IPC shards for cp_rank={cp_rank} disagree on " + f"sequence range for batch_item_id={batch_item_id}." + ) + vocab_range = (vocab_start, vocab_end) + prior_vocab_range = vocab_ranges.setdefault(tp_rank, vocab_range) + if prior_vocab_range != vocab_range: + raise ValueError( + f"Dense teacher IPC shards for tp_rank={tp_rank} disagree on " + f"vocabulary range for batch_item_id={batch_item_id}." + ) + + compact_geometry = validate_compact_teacher_ipc_handle(shard) + if compact_geometry is not None: + token_offset, stored_seq_len, used_tokens, _ = compact_geometry + valid_seq_len = int(shard["valid_seq_len"]) + valid_seq_lengths.add(valid_seq_len) + storage_capacity_tokens = int(shard["storage_capacity_tokens"]) + geometry = ( + token_offset, + stored_seq_len, + used_tokens, + storage_capacity_tokens, + ) + prior_geometry = compact_geometry_by_cp.setdefault(cp_rank, geometry) + if prior_geometry != geometry: + raise ValueError( + "Compact dense teacher IPC TP shards disagree on physical " + f"storage geometry for cp_rank={cp_rank}, " + f"batch_item_id={batch_item_id}." + ) + + if len(layouts) != 1: + raise ValueError( + "Dense teacher IPC shards mix incompatible storage layouts for " + f"batch_item_id={batch_item_id}: {sorted(map(str, layouts))}." + ) + if layouts != {None} and len(valid_seq_lengths) != 1: + raise ValueError( + "Compact dense teacher IPC shards disagree on valid_seq_len for " + f"batch_item_id={batch_item_id}: {sorted(valid_seq_lengths)}." + ) + + expected_coordinates = { + (tp_rank, cp_rank) for tp_rank in range(tp_size) for cp_rank in range(cp_size) + } + if coordinates != expected_coordinates: + raise ValueError( + "Dense teacher IPC has incomplete TP/CP shard coverage for " + f"batch_item_id={batch_item_id}: missing=" + f"{sorted(expected_coordinates - coordinates)}, unexpected=" + f"{sorted(coordinates - expected_coordinates)}." + ) + + sequence_cursor = 0 + for cp_rank in range(cp_size): + seq_start, seq_end = sequence_ranges[cp_rank] + if seq_start != sequence_cursor: + raise ValueError( + "Dense teacher IPC sequence shards contain a gap or overlap for " + f"batch_item_id={batch_item_id} before cp_rank={cp_rank}: " + f"expected start {sequence_cursor}, got {seq_start}." + ) + sequence_cursor = seq_end + if sequence_cursor != full_seq_len: + raise ValueError( + "Dense teacher IPC sequence shards do not cover the logical axis for " + f"batch_item_id={batch_item_id}: covered={sequence_cursor}, " + f"expected={full_seq_len}." + ) + + vocab_cursor = 0 + for tp_rank in range(tp_size): + vocab_start, vocab_end = vocab_ranges[tp_rank] + if vocab_start != vocab_cursor: + raise ValueError( + "Dense teacher IPC vocabulary shards contain a gap or overlap for " + f"batch_item_id={batch_item_id} before tp_rank={tp_rank}: " + f"expected start {vocab_cursor}, got {vocab_start}." + ) + vocab_cursor = vocab_end + if vocab_cursor != full_vocab_size: + raise ValueError( + "Dense teacher IPC vocabulary shards do not cover the logical axis for " + f"batch_item_id={batch_item_id}: covered={vocab_cursor}, " + f"expected={full_vocab_size}." + ) + + def aggregate_per_sample_handles( worker_results: list[dict[str, Any]], + canonical_batch_item_ids: Optional[Iterable[int]] = None, ) -> list[dict[str, Any]]: """Flatten teacher per-sample IPC handles into a global-batch-ordered list. Each worker returns ``{"dp_rank": int, "per_sample_handles": list}`` where the handle list is in local sample order; the several workers sharing a - ``dp_rank`` are TP/CP replicas that each contribute one shard per sample. + ``dp_rank`` are TP/CP/PP replicas that each contribute one shard per sample. Concatenating samples in ``sorted(dp_rank)`` order reproduces the original global sample order (rank 0 holds the first ``gbs/dp`` samples, rank 1 the next, ...), so the result is a length-``gbs`` list independent of the teacher's DP degree. Element ``i`` is ``{"teacher_shards": [shard, ...]}`` holding all TP×CP shards of global sample ``i``. + + Workers that contribute no handles are skipped. Under pipeline parallelism + only the LAST stage holds logits, so every earlier stage returns an empty + list; without this filter the per-``dp_rank`` length check below would fire + on those empty contributions. + + When ``canonical_batch_item_ids`` is supplied, positional ordering is not + used: every record must carry ``batch_item_id``, TP/CP replicas on one DP + rank must report identical identity sets, and every sample must have a + complete, non-overlapping TP×CP shard grid covering its full logical + sequence/vocabulary rectangle. The aggregate is returned in exact canonical + order. Omitting the IDs retains the legacy unpacked behavior. """ + nonempty_results = [ + worker_result + for worker_result in worker_results + if worker_result["per_sample_handles"] + ] + assert nonempty_results, ( + "No teacher worker returned full-logits IPC handles. Under pipeline " + "parallelism only the last stage produces them, so this means the " + "last-stage results were not collected." + ) + if canonical_batch_item_ids is not None: + canonical_ids = tuple(int(value) for value in canonical_batch_item_ids) + if len(set(canonical_ids)) != len(canonical_ids): + raise ValueError( + f"canonical_batch_item_ids must be unique, got {canonical_ids}." + ) + shards_by_item: dict[int, list[dict[str, Any]]] = {} + owner_dp_by_item: dict[int, int] = {} + identity_sets_by_dp: dict[int, list[set[int]]] = {} + for worker_result in nonempty_results: + dp_rank = int(worker_result["dp_rank"]) + seen_by_worker: set[int] = set() + for handle in worker_result["per_sample_handles"]: + if "batch_item_id" not in handle: + raise ValueError( + "Identity-aware dense teacher IPC requires every handle " + "record to carry batch_item_id." + ) + batch_item_id = int(handle["batch_item_id"]) + if batch_item_id in seen_by_worker: + raise ValueError( + "A teacher worker returned duplicate handles for " + f"batch_item_id={batch_item_id}." + ) + seen_by_worker.add(batch_item_id) + owner_dp = owner_dp_by_item.setdefault(batch_item_id, dp_rank) + if owner_dp != dp_rank: + raise ValueError( + f"batch_item_id={batch_item_id} was returned by multiple " + f"data-parallel ranks ({owner_dp} and {dp_rank})." + ) + shards_by_item.setdefault(batch_item_id, []).append(handle) + identity_sets_by_dp.setdefault(dp_rank, []).append(seen_by_worker) + for dp_rank, worker_identity_sets in identity_sets_by_dp.items(): + expected = worker_identity_sets[0] + if any( + identity_set != expected for identity_set in worker_identity_sets[1:] + ): + raise ValueError( + f"Teacher shard replicas for dp_rank={dp_rank} returned " + f"inconsistent batch_item_id sets: {worker_identity_sets}." + ) + actual_ids = set(shards_by_item) + canonical_id_set = set(canonical_ids) + if actual_ids != canonical_id_set: + missing = sorted(canonical_id_set - actual_ids) + unexpected = sorted(actual_ids - canonical_id_set) + raise ValueError( + "Dense teacher IPC identities do not cover the canonical batch: " + f"missing={missing}, unexpected={unexpected}." + ) + for batch_item_id in canonical_ids: + _validate_dense_ipc_shard_coverage( + batch_item_id, shards_by_item[batch_item_id] + ) + return [ + { + "batch_item_id": batch_item_id, + "teacher_shards": shards_by_item[batch_item_id], + } + for batch_item_id in canonical_ids + ] + handles_by_dp_rank: dict[int, list[list[dict[str, Any]]]] = {} - for worker_result in worker_results: + for worker_result in nonempty_results: + per_sample_handles = worker_result["per_sample_handles"] dp_rank = worker_result["dp_rank"] - handles_by_dp_rank.setdefault(dp_rank, []).append( - worker_result["per_sample_handles"] - ) + handles_by_dp_rank.setdefault(dp_rank, []).append(per_sample_handles) aggregated: list[dict[str, Any]] = [] for dp_rank in sorted(handles_by_dp_rank.keys()): worker_handles_in_dp = handles_by_dp_rank[dp_rank] diff --git a/nemo_rl/models/policy/workers/automodel_policy_worker.py b/nemo_rl/models/policy/workers/automodel_policy_worker.py index 48e26b8b756..4f1130d9c87 100644 --- a/nemo_rl/models/policy/workers/automodel_policy_worker.py +++ b/nemo_rl/models/policy/workers/automodel_policy_worker.py @@ -31,6 +31,7 @@ from nemo_rl.algorithms.logits_sampling_utils import ( TrainingSamplingParams, ) +from nemo_rl.algorithms.loss.csd_packing import CSD_LOGICAL_METRICS_KEY from nemo_rl.algorithms.loss.interfaces import LossFunction from nemo_rl.algorithms.metric_utils import LEARNING_RATE_KEY from nemo_rl.data_plane.worker_mixin import TQWorkerMixin @@ -59,6 +60,7 @@ LogprobsPostProcessor, LossPostProcessor, ScorePostProcessor, + SparseLogitsPostProcessor, TopkLogitsPostProcessor, aggregate_training_statistics, automodel_forward_backward, @@ -75,6 +77,8 @@ from nemo_rl.models.policy.utils import ( enable_deterministic_algorithms, ensure_teacher_ipc_buffer, + extract_batch_item_ids, + get_handle_from_tensor, get_runtime_env_for_policy_worker, ) from nemo_rl.models.policy.workers.base_policy_worker import AbstractPolicyWorker @@ -285,12 +289,14 @@ def __init__( # IPC-imported storage without pinning an orphaned producer allocation. self._teacher_ipc_storage: Optional[torch.Tensor] = None self._teacher_ipc_handle: Optional[tuple[Any, ...]] = None + self._teacher_sparse_ipc_buffer: Optional[list[torch.Tensor]] = None # Validate configuration and prepare runtime settings runtime_config = validate_and_prepare_config( config=config, processor=self.processor, rank=0, # Temporary, will be updated after distributed init + init_optimizer=init_optimizer, ) # Set up distributed environment (returns DistributedContext) @@ -408,6 +414,7 @@ def __init__( self.is_generation_colocated, self.sampling_params, _runtime_is_reward_model, # Duplicate, already set as _is_reward_model + self.model_load_dtype, ) = runtime_config def _update_moe_gate_bias_if_supported(self) -> None: @@ -504,6 +511,18 @@ def on_microbatch_start(mb_idx): batch = gb_result["batch"] global_valid_seqs = gb_result["global_valid_seqs"] global_valid_toks = gb_result["global_valid_toks"] + global_valid_kd_toks = gb_result.get( + "global_valid_kd_toks", global_valid_toks + ) + global_valid_chunks_by_idx = gb_result["global_valid_chunks_by_idx"] + # CSD gathers the full sequence and computes a replicated + # chunk numerator on every CP rank. Its denominator therefore + # includes the CP copies; FSDP's DP x CP reduction sums them back + # to one natural global-batch chunk mean. + global_valid_chunks_by_idx = { + i: count * self.cp_size + for i, count in global_valid_chunks_by_idx.items() + } self.optimizer.zero_grad() @@ -529,6 +548,8 @@ def on_microbatch_start(mb_idx): allow_flash_attn_args=self.allow_flash_attn_args, global_valid_seqs=global_valid_seqs, global_valid_toks=global_valid_toks, + global_valid_kd_toks=global_valid_kd_toks, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, sampling_params=self.sampling_params, sequence_dim=sequence_dim, dp_size=self.dp_size, @@ -544,19 +565,34 @@ def on_microbatch_start(mb_idx): for mb_idx, (loss, loss_metrics) in enumerate(mb_results): # Only process valid (non-dummy) batches for metrics if mb_idx < iterator_len: - num_valid_samples = loss_metrics["num_valid_samples"] - loss_metrics[LEARNING_RATE_KEY] = self.optimizer.param_groups[ - 0 - ]["lr"] - loss_metrics["global_valid_seqs"] = global_valid_seqs.item() - loss_metrics["global_valid_toks"] = global_valid_toks.item() - - if num_valid_samples > 0: - # Metrics were materialized together by the loss; - # undo this worker's per-global-batch scaling without - # synchronizing the loss tensor again. - mb_losses.append(loss_metrics["loss"] * num_global_batches) - all_mb_metrics.append(loss_metrics) + logical_metrics = loss_metrics.pop( + CSD_LOGICAL_METRICS_KEY, None + ) + metric_records = ( + logical_metrics + if logical_metrics is not None + else [loss_metrics] + ) + valid_metric_records = [] + for metric_record in metric_records: + if metric_record["num_valid_samples"] <= 0: + continue + metric_record[LEARNING_RATE_KEY] = ( + self.optimizer.param_groups[0]["lr"] + ) + metric_record["global_valid_seqs"] = ( + global_valid_seqs.item() + ) + metric_record["global_valid_toks"] = ( + global_valid_toks.item() + ) + valid_metric_records.append(metric_record) + if valid_metric_records: + mb_losses.append( + sum(record["loss"] for record in valid_metric_records) + * num_global_batches + ) + all_mb_metrics.extend(valid_metric_records) grad_norm: Optional[float | torch.Tensor] = None if not eval_mode: @@ -903,6 +939,17 @@ def get_full_logits_ipc( full_seq_len = target_local_seq * self.cp_size global_seq_start = cp_rank * full_seq_len // self.cp_size + planned_batch_capacity = forward_batch_size + if data.micro_batch_indices is not None: + if len(data.micro_batch_indices) != 1: + raise ValueError( + "Dense teacher IPC expects one local microbatch plan, got " + f"{len(data.micro_batch_indices)}." + ) + planned_batch_capacity = max( + end - start for start, end in data.micro_batch_indices[0] + ) + per_sample_handles: list[dict[str, Any]] = [] storage: Optional[torch.Tensor] = None payload_ipc: Optional[tuple[Any, ...]] = None @@ -944,21 +991,36 @@ def get_full_logits_ipc( vals, (0, 0, 0, pad_needed, 0, 0), mode="constant", value=0.0 ) batch_size_mb, seq_len_mb, local_vocab_size = vals.shape + batch_item_ids = extract_batch_item_ids( + processed_mb.data_dict, + batch_size_mb, + required=self.enable_seq_packing, + ) - self._teacher_ipc_storage, self._teacher_ipc_handle = ( - ensure_teacher_ipc_buffer( - self._teacher_ipc_storage, - self._teacher_ipc_handle, - iterator_len, - batch_size_mb, - target_local_seq, - local_vocab_size, - vals.dtype, - vals.device, + # Allocate against the complete planned cardinality before the + # first handle is published. Later variable-cardinality bins + # must never grow the storage and invalidate earlier handles. + if storage is None: + self._teacher_ipc_storage, self._teacher_ipc_handle = ( + ensure_teacher_ipc_buffer( + self._teacher_ipc_storage, + self._teacher_ipc_handle, + iterator_len, + planned_batch_capacity, + target_local_seq, + local_vocab_size, + vals.dtype, + vals.device, + ) + ) + storage = self._teacher_ipc_storage + payload_ipc = self._teacher_ipc_handle + assert storage is not None and payload_ipc is not None + if batch_size_mb > storage.shape[1]: + raise ValueError( + "Dense teacher IPC microbatch exceeds its preallocated " + f"cardinality: rows={batch_size_mb}, capacity={storage.shape[1]}." ) - ) - storage = self._teacher_ipc_storage - payload_ipc = self._teacher_ipc_handle storage[buf_idx, :batch_size_mb, :seq_len_mb, :local_vocab_size].copy_( vals ) @@ -967,28 +1029,31 @@ def get_full_logits_ipc( vocab_start_index = tp_rank * local_vocab_size vocab_end_index = (tp_rank + 1) * local_vocab_size for sample_index_in_buf in range(batch_size_mb): - per_sample_handles.append( - { - "payload_ipc": payload_ipc, - "buf_idx": buf_idx, - "sample_index_in_buf": sample_index_in_buf, - "storage_shape": tuple(storage.shape), - "actual_shape": (target_local_seq, local_vocab_size), - "dtype": storage.dtype, - "tp_rank": tp_rank, - "cp_rank": cp_rank, - "tp_size": self.tp_size, - "cp_size": self.cp_size, - "world_rank": world_rank, - "vocab_start_index": vocab_start_index, - "vocab_end_index": vocab_end_index, - "global_seq_start": global_seq_start, - "full_vocab_size": full_vocab_size, - "full_seq_len": full_seq_len, - "vocab_sharded": self.tp_size > 1, - "sequence_sharded": self.cp_size > 1, - } - ) + handle_record = { + "payload_ipc": payload_ipc, + "buf_idx": buf_idx, + "sample_index_in_buf": sample_index_in_buf, + "storage_shape": tuple(storage.shape), + "actual_shape": (target_local_seq, local_vocab_size), + "dtype": storage.dtype, + "tp_rank": tp_rank, + "cp_rank": cp_rank, + "tp_size": self.tp_size, + "cp_size": self.cp_size, + "world_rank": world_rank, + "vocab_start_index": vocab_start_index, + "vocab_end_index": vocab_end_index, + "global_seq_start": global_seq_start, + "full_vocab_size": full_vocab_size, + "full_seq_len": full_seq_len, + "vocab_sharded": self.tp_size > 1, + "sequence_sharded": self.cp_size > 1, + } + if batch_item_ids is not None: + handle_record["batch_item_id"] = batch_item_ids[ + sample_index_in_buf + ] + per_sample_handles.append(handle_record) # The storage copies above are async on the current stream; force them # to complete before the IPC handles are consumed by the student # process, so the consumer can't observe a partially written buffer @@ -996,10 +1061,228 @@ def get_full_logits_ipc( torch.cuda.synchronize() return {"per_sample_handles": per_sample_handles, "dp_rank": dp_rank} + @wrap_with_nvtx_name("automodel_policy_worker/get_topk_logits_ipc") + def get_topk_logits_ipc( + self, + data: BatchedDataDict[Any], + *, + k: int, + temperature: float, + vocab_size: Optional[int] = None, + micro_batch_size: Optional[int] = None, + support_mode: str, + gt_filter_topk: Optional[int] = None, + ) -> dict[str, Any]: + """Run a teacher forward and expose row-top-k + logZ via CUDA IPC.""" + if k <= 0: + raise ValueError(f"k must be > 0 for get_topk_logits_ipc, got {k}") + if temperature <= 0.0: + raise ValueError( + f"temperature must be > 0 for get_topk_logits_ipc, got {temperature}" + ) + if support_mode != "row_topk": + raise NotImplementedError( + "Automodel sparse teacher IPC currently supports only " + f"support_mode='row_topk', got {support_mode!r}." + ) + gt_filter_topk_value = int(gt_filter_topk or 0) + if gt_filter_topk_value < 0: + raise ValueError( + "gt_filter_topk must be >= 0 for get_topk_logits_ipc, " + f"got {gt_filter_topk_value}" + ) + + forward_batch_size = ( + micro_batch_size + if micro_batch_size is not None + else self.cfg["logprob_batch_size"] + ) + sequence_dim, seq_dim_size = check_sequence_dim(data) + target_local_seq = seq_dim_size // self.cp_size + post_processor = SparseLogitsPostProcessor( + cfg=self.cfg, + tp_mesh=self.tp_mesh, + cp_mesh=self.cp_mesh, + cp_size=self.cp_size, + k=k, + temperature=temperature, + vocab_size=vocab_size, + gt_filter_topk=gt_filter_topk_value, + ) + + output_logits: list[torch.Tensor] = [] + output_indices: list[torch.Tensor] = [] + output_log_z: list[torch.Tensor] = [] + output_gt_in_topk: list[torch.Tensor] = [] + self.model.eval() + with torch.no_grad(): + data.to("cuda") + processed_iterator, iterator_len = get_microbatch_iterator( + data, + self.cfg, + forward_batch_size, + self.dp_mesh, + tokenizer=self.tokenizer, + ) + for batch_idx, processed_mb in enumerate(processed_iterator): + processed_inputs = processed_mb.processed_inputs + prepared = prepare_model_forward( + self.model, + processed_inputs, + device_mesh=self.device_mesh, + cp_size=self.cp_size, + padding_token_id=self.tokenizer.pad_token_id or 0, + is_reward_model=False, + allow_flash_attn_args=self.allow_flash_attn_args, + ) + with prepared.model_context_factory(), self._autocast_context(): + sparse_values, _metrics, _ = forward_with_post_processing_fn( + model=self.model, + prepared=prepared, + post_processing_fn=post_processor, + processed_mb=processed_mb, + sampling_params=self.sampling_params, + sequence_dim=sequence_dim, + ) + if batch_idx >= iterator_len: + continue + topk_logits, topk_indices, log_z, gt_in_topk = sparse_values + pad_needed = target_local_seq - topk_logits.shape[1] + if pad_needed > 0: + topk_logits = torch.nn.functional.pad( + topk_logits, + (0, 0, 0, pad_needed, 0, 0), + mode="constant", + value=0.0, + ) + topk_indices = torch.nn.functional.pad( + topk_indices, + (0, 0, 0, pad_needed, 0, 0), + mode="constant", + value=0, + ) + log_z = torch.nn.functional.pad( + log_z, + (0, pad_needed, 0, 0), + mode="constant", + value=0.0, + ) + if gt_in_topk is not None: + gt_in_topk = torch.nn.functional.pad( + gt_in_topk, + (0, pad_needed, 0, 0), + mode="constant", + value=False, + ) + output_logits.append(topk_logits.contiguous()) + output_indices.append(topk_indices.to(torch.int32).contiguous()) + output_log_z.append(log_z.contiguous()) + if gt_in_topk is not None: + output_gt_in_topk.append(gt_in_topk.to(torch.int32).contiguous()) + + if not output_logits: + raise RuntimeError("No valid microbatches produced sparse teacher logits.") + final_logits = torch.cat(output_logits, dim=0) + final_indices = torch.cat(output_indices, dim=0) + final_log_z = torch.cat(output_log_z, dim=0) + final_gt_in_topk = ( + torch.cat(output_gt_in_topk, dim=0) if output_gt_in_topk else None + ) + + sparse_sources = [final_logits, final_indices, final_log_z] + if final_gt_in_topk is not None: + sparse_sources.append(final_gt_in_topk) + existing_sparse_storage = self._teacher_sparse_ipc_buffer or [] + sparse_storage: list[torch.Tensor] = [] + sparse_views: list[torch.Tensor] = [] + for tensor_idx, source in enumerate(sparse_sources): + storage = ( + existing_sparse_storage[tensor_idx] + if tensor_idx < len(existing_sparse_storage) + else None + ) + needs_reallocation = ( + storage is None + or storage.ndim != source.ndim + or any( + storage.shape[dim] < source.shape[dim] for dim in range(source.ndim) + ) + or storage.dtype != source.dtype + or storage.device != source.device + ) + if needs_reallocation: + storage = torch.empty_like(source) + target = storage[tuple(slice(0, size) for size in source.shape)] + target.copy_(source) + sparse_storage.append(storage) + sparse_views.append(target) + self._teacher_sparse_ipc_buffer = sparse_storage + final_logits = sparse_views[0] + final_indices = sparse_views[1] + final_log_z = sparse_views[2] + final_gt_in_topk = sparse_views[3] if len(sparse_views) == 4 else None + torch.cuda.synchronize() + + sparse_bytes = ( + final_logits.numel() * final_logits.element_size() + + final_indices.numel() * final_indices.element_size() + + final_log_z.numel() * final_log_z.element_size() + ) + if final_gt_in_topk is not None: + sparse_bytes += final_gt_in_topk.numel() * final_gt_in_topk.element_size() + full_vocab = int(vocab_size or final_indices.max().item() + 1) + full_bytes_estimate = ( + final_logits.shape[0] + * final_logits.shape[1] + * full_vocab + * final_logits.element_size() + ) + print( + "Sparse teacher IPC buffer: " + f"shape={tuple(final_logits.shape)} " + f"support_mode={support_mode} " + f"bytes={sparse_bytes / (1024**3):.2f}GiB " + f"full_est={full_bytes_estimate / (1024**3):.2f}GiB", + flush=True, + ) + + cp_rank = self.cp_mesh.get_local_rank() + dp_rank = self.dp_mesh.get_local_rank() + global_seq_start = cp_rank * final_logits.shape[1] + per_sample_handles: list[dict[str, Any]] = [] + for sample_idx in range(final_logits.shape[0]): + logits_sample = final_logits[sample_idx] + indices_sample = final_indices[sample_idx] + log_z_sample = final_log_z[sample_idx] + handle: dict[str, Any] = { + "topk_logits_ipc": get_handle_from_tensor(logits_sample), + "topk_indices_ipc": get_handle_from_tensor(indices_sample), + "log_z_ipc": get_handle_from_tensor(log_z_sample), + "topk_shape": tuple(logits_sample.shape), + "topk_dtype": logits_sample.dtype, + "indices_dtype": indices_sample.dtype, + "log_z_dtype": log_z_sample.dtype, + "cp_rank": cp_rank, + "cp_size": self.cp_size, + "global_seq_start": global_seq_start, + "full_seq_len": final_logits.shape[1] * self.cp_size, + } + if final_gt_in_topk is not None: + gt_sample = final_gt_in_topk[sample_idx] + handle.update( + { + "gt_in_topk_ipc": get_handle_from_tensor(gt_sample), + "gt_in_topk_dtype": gt_sample.dtype, + } + ) + per_sample_handles.append(handle) + return {"per_sample_handles": per_sample_handles, "dp_rank": dp_rank} + def release_ipc_buffer(self) -> None: """Free the persistent teacher-logit IPC storage. Called once at end of training/validation.""" self._teacher_ipc_storage = None self._teacher_ipc_handle = None + self._teacher_sparse_ipc_buffer = None gc.collect() torch.cuda.empty_cache() diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 274d332e4d3..a1abb2f3235 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -57,6 +57,7 @@ from nemo_rl.algorithms.logits_sampling_utils import ( TrainingSamplingParams, ) +from nemo_rl.algorithms.loss.csd_packing import CSD_LOGICAL_METRICS_KEY from nemo_rl.algorithms.loss.interfaces import LossFunction from nemo_rl.algorithms.loss.loss_functions import ClippedPGLossFn from nemo_rl.algorithms.loss.utils import rescale_loss_metrics @@ -68,6 +69,7 @@ ) from nemo_rl.data_plane.worker_mixin import TQWorkerMixin from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.model_utils import cp_load_balanced_segments from nemo_rl.distributed.named_sharding import NamedSharding from nemo_rl.models.generation.interfaces import GenerationDatumSpec from nemo_rl.models.generation.megatron.megatron_worker import ( @@ -122,11 +124,15 @@ validate_model_paths, ) from nemo_rl.models.megatron.train import ( + FullLogitsPostProcessor, LogprobsPostProcessor, LossPostProcessor, + SparseTeacherLogitsPostProcessor, + StreamedFullLogitsMetadata, TeacherFullPayloadPostProcessor, TopkLogitsPostProcessor, aggregate_training_statistics, + loss_returns_context_parallel_partial, megatron_forward_backward, should_reduce_loss_across_context_parallel, strip_context_parallel_local_loss_metric, @@ -139,10 +145,20 @@ TeacherFullPayloadOutputSpec, ) from nemo_rl.models.policy.utils import ( + DENSE_TEACHER_IPC_FLAT_LAYOUT, broadcast_hf_buckets_via_distributed_impl, + can_reuse_teacher_ipc_token_buffer, connect_rollout_engines_from_distributed, disconnect_rollout_engines_from_distributed, + ensure_teacher_ipc_buffer, + ensure_teacher_ipc_token_buffer, + extract_batch_item_ids, + extract_teacher_ipc_valid_lengths, + get_dense_ipc_sequence_layout, + get_handle_from_tensor, get_runtime_env_for_policy_worker, + localize_teacher_ipc_valid_lengths, + partition_teacher_ipc_token_buffer, send_hf_buckets_via_ipc_actor_impl, ) from nemo_rl.models.policy.workers.base_policy_worker import AbstractPolicyWorker @@ -631,6 +647,43 @@ class MegatronPolicyWorkerImpl( _remote_sparse_refit: Any = None _async_checkpoint_cuda_cache_active: bool = False + def _ensure_teacher_ipc_storage( + self, + *, + total_tokens: int, + vocab_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Reuse or replace the compact IPC slab without overlapping owners.""" + storage = self._teacher_ipc_storage + if storage is not None and not can_reuse_teacher_ipc_token_buffer( + storage, + total_tokens=total_tokens, + vocab_size=vocab_size, + dtype=dtype, + device=device, + ): + # A new get_full_logits_ipc call starts only after the previous + # student train call has returned, so its imported views are no + # longer live. Drop the producer owner and stale handle before the + # replacement allocation; otherwise both multi-GiB slabs overlap. + self._teacher_ipc_storage = None + self._teacher_ipc_handles = [] + del storage + gc.collect() + torch.cuda.empty_cache() + + storage = ensure_teacher_ipc_token_buffer( + self._teacher_ipc_storage, + total_tokens=total_tokens, + vocab_size=vocab_size, + dtype=dtype, + device=device, + ) + self._teacher_ipc_storage = storage + return storage + def __repr__(self): """Customizes the actor's prefix in the Ray logs. @@ -793,6 +846,19 @@ def __init__( init_telemetry_worker() self.cfg = config + # Persistent compact CUDA IPC token slab for cross-tokenizer teacher + # full logits, lazily allocated on the first ``get_full_logits_ipc`` + # call (grown as needed), reused across calls, and freed via + # ``release_ipc_buffer``. Logical row and TP/CP geometry remain in each + # handle while zero padding is omitted from the physical allocation. + self._teacher_ipc_storage: Optional[torch.Tensor] = None + self._teacher_ipc_handles: list[tuple[Any, ...]] = [] + # Persistent ``[N_mb, B, T_local, width]`` slabs for sparse teacher + # rows, one per exported field, with their IPC handles cached across + # calls; allocated by ``get_sparse_logits_ipc`` and freed by + # ``release_ipc_buffer``. + self._teacher_sparse_ipc_storage: dict[str, torch.Tensor] = {} + self._teacher_sparse_ipc_handles: dict[str, tuple[Any, ...]] = {} self._router_replay_enabled = router_replay_enabled(config) # HF param names to MXFP8-quantize on the trainer during refit; set via # enable_refit_prequantize() when vllm_cfg.refit_prequantize is on. @@ -1265,15 +1331,12 @@ def train( ) -> dict[str, Any]: """Train the policy on a batch of data with a given loss function. - ``check_dim_skip_keys`` is accepted for parity with the DTensor - worker (cross-tokenizer ride-along tensors whose dim 1 is not the - student sequence axis). Megatron doesn't run cross-tokenizer, so it - must be None. + ``check_dim_skip_keys`` names cross-tokenizer ride-along tensors whose + dim 1 is NOT the student sequence axis (teacher tokenization / alignment + keys); it is forwarded to the microbatch iterator's sequence-dim + validation so those keys pass through untouched, matching the Automodel + worker. """ - assert check_dim_skip_keys is None, ( - "check_dim_skip_keys is only supported by the v2 DTensor worker; " - "Megatron does not run cross-tokenizer distillation." - ) self.timer.start("train") # Note: zero_grad_buffer is called at the start of each global batch iteration # in the loop below, so we don't need to call it here. @@ -1336,6 +1399,10 @@ def train( batch = gb_result["batch"] global_valid_seqs = gb_result["global_valid_seqs"] global_valid_toks = gb_result["global_valid_toks"] + global_valid_kd_toks = gb_result.get( + "global_valid_kd_toks", global_valid_toks + ) + global_valid_chunks_by_idx = gb_result["global_valid_chunks_by_idx"] # Pre-compute the MTP loss mask, only when MTP is enabled, so # process_microbatch can pack it. @@ -1368,6 +1435,7 @@ def train( mtp_enabled=self.mtp_enabled, create_router_padding_mask=self.create_router_padding_mask, allow_unmasked_chunkwise_cp=self.allow_unmasked_chunkwise_cp, + skip_keys=check_dim_skip_keys, ) # Track total microbatches for MoE aux-loss averaging total_num_microbatches += int(num_microbatches) @@ -1376,6 +1444,7 @@ def train( loss_fn=loss_fn, cfg=self.cfg, num_microbatches=num_microbatches, + cp_normalize=not loss_returns_context_parallel_partial(loss_fn), sampling_params=self.sampling_params, draft_model=self.draft_model, teacher_output_layer_weight_by_index=self._opd_full_teacher_lm_heads, @@ -1429,6 +1498,8 @@ def train( defer_fp32_logits=self.defer_fp32_logits, global_valid_seqs=global_valid_seqs, global_valid_toks=global_valid_toks, + global_valid_kd_toks=global_valid_kd_toks, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, sampling_params=self.sampling_params, straggler_timer=self.mcore_state.straggler_timer, model_slices_context_parallel_inputs=self.model_slices_context_parallel_inputs, @@ -1528,20 +1599,33 @@ def train( gb_loss_metrics = [] mb_losses = [] for x in losses_reduced: - loss_metrics = {} - for k in x.keys(): - if "_min" in k or "_max" in k: - loss_metrics[k] = x[k] - else: - loss_metrics[k] = x[k] / num_global_batches - gb_loss_metrics.append(loss_metrics) - curr_lr = self.scheduler.get_lr(self.optimizer.param_groups[0]) - curr_wd = self.scheduler.get_wd() - loss_metrics[LEARNING_RATE_KEY] = curr_lr - loss_metrics["wd"] = curr_wd - loss_metrics["global_valid_seqs"] = global_valid_seqs.detach() - loss_metrics["global_valid_toks"] = global_valid_toks.detach() - mb_losses.append(loss_metrics["loss"]) + logical_metrics = x.get(CSD_LOGICAL_METRICS_KEY) + metric_records = ( + logical_metrics if logical_metrics is not None else [x] + ) + for metric_record in metric_records: + if metric_record["num_valid_samples"] <= 0: + continue + loss_metrics = {} + for key, value in metric_record.items(): + if "_min" in key or "_max" in key: + loss_metrics[key] = value + else: + loss_metrics[key] = value / num_global_batches + curr_lr = self.scheduler.get_lr( + self.optimizer.param_groups[0] + ) + curr_wd = self.scheduler.get_wd() + loss_metrics[LEARNING_RATE_KEY] = curr_lr + loss_metrics["wd"] = curr_wd + loss_metrics["global_valid_seqs"] = ( + global_valid_seqs.detach() + ) + loss_metrics["global_valid_toks"] = ( + global_valid_toks.detach() + ) + gb_loss_metrics.append(loss_metrics) + mb_losses.append(loss_metrics["loss"]) else: gb_loss_metrics = None @@ -2075,6 +2159,7 @@ def _train_microbatch_body( loss_fn=loss_fn, cfg=self.cfg, num_microbatches=num_microbatches, + cp_normalize=not loss_returns_context_parallel_partial(loss_fn), sampling_params=self.sampling_params, draft_model=self.draft_model, defer_draft_normalization=True, @@ -3138,6 +3223,706 @@ def get_topk_logits( [{"topk_logits": topk_logits.cpu(), "topk_indices": topk_indices.cpu()}] ) + @torch.no_grad() + @wrap_with_nvtx_name("megatron_policy_worker/get_full_logits_ipc") + def get_full_logits_ipc( + self, + data: BatchedDataDict[GenerationDatumSpec], + micro_batch_size: Optional[int] = None, + ) -> dict[str, Any]: + """Teacher forward; full-vocab logits exposed via persistent CUDA IPC storage. + + Megatron counterpart of ``AutomodelPolicyWorker.get_full_logits_ipc`` for + cross-tokenizer distillation. Each microbatch's raw logit shard is + written into a compact ``[sum(valid_local_T), V]`` persistent token + slab. Packed rows store only the overlap of their authoritative raw + ``input_lengths`` prefix with this teacher CP rank; the postprocessor + already zeroes the omitted logical suffix. Unpacked rows retain their + complete width because their padding logits are not guaranteed zero. + The handle preserves ``buf_idx=0`` / bin-local + ``sample_index_in_buf`` identity while adding physical token offsets. + Returns ``{"per_sample_handles": list, "dp_rank": int}`` + with TP/CP shard metadata (``vocab_start_index``, ``global_seq_start``, + ...) used to reassemble the global ``[T_t, V_t]`` teacher logits. + + Tensor, context and pipeline parallelism are supported for models using + NeMo-RL's generic Megatron batch path and may differ from the student's. + Models that pack and context-shard inside their own forward must use + ``context_parallel_size == 1`` for this unpacked IPC path. Under PP only + the LAST stage runs the output layer, so ``megatron_forward_backward`` + returns an empty output list on every earlier stage and this method + returns no handles there; :func:`aggregate_per_sample_handles` skips + those empty contributions. + Multi-node PP is guarded separately by + :func:`assert_xtoken_ipc_node_local`, which requires the whole + ``tp * cp * pp`` group to fit inside one node so the producing stage is + co-located with the importing student ranks. + + Sequence-packed outputs are reconstructed one physical sample at a time + by :class:`FullLogitsPostProcessor` and streamed directly into the + persistent IPC slab. + """ + if ( + self.delegate_pack_to_model + and parallel_state.get_context_parallel_world_size() > 1 + ): + raise NotImplementedError( + "get_full_logits_ipc does not support context parallelism for " + "models that pack and context-shard inside their own forward. " + "Set context_parallel_size=1 for this teacher." + ) + + forward_batch_size = ( + micro_batch_size + if micro_batch_size is not None + else self.cfg["logprob_batch_size"] + ) + self.model.eval() + + # Size the persistent slab from logical input before the forward. Packed + # output streams into this allocation microbatch-by-microbatch, so it + # must never fall back to a post-forward allocation that overlaps model + # outputs. Pipeline stages without the output layer legitimately emit + # no handles and do not allocate a slab. + cp_rank = parallel_state.get_context_parallel_rank() + cp_size = parallel_state.get_context_parallel_world_size() + tp_size = parallel_state.get_tensor_model_parallel_world_size() + full_seq_len, target_local_seq, global_seq_start = ( + get_dense_ipc_sequence_layout(data, cp_rank=cp_rank, cp_size=cp_size) + ) + planned_valid_lengths = extract_teacher_ipc_valid_lengths( + data, + data.size, + full_seq_len=full_seq_len, + compact_padding=self.cfg["sequence_packing"]["enabled"], + ) + planned_total_stored_tokens = sum( + localize_teacher_ipc_valid_lengths( + planned_valid_lengths, + full_seq_len=full_seq_len, + cp_rank=cp_rank, + cp_size=cp_size, + ) + ) + pack = self.cfg["sequence_packing"]["enabled"] + stream_packed_outputs = pack and parallel_state.is_pipeline_last_stage( + ignore_virtual=True + ) + packed_output_storage: Optional[torch.Tensor] = None + planned_local_vocab_size: Optional[int] = None + if stream_packed_outputs: + final_padded_vocab_size = int(self.final_padded_vocab_size) + if ( + tp_size <= 0 + or final_padded_vocab_size <= 0 + or final_padded_vocab_size % tp_size != 0 + ): + raise ValueError( + "The padded teacher vocabulary must be positive and divisible " + "by tensor parallelism for packed IPC streaming: " + f"vocab={final_padded_vocab_size}, tp={tp_size}." + ) + planned_local_vocab_size = final_padded_vocab_size // tp_size + current_storage = self._teacher_ipc_storage + storage_device = ( + current_storage.device + if current_storage is not None + else torch.device("cuda", torch.cuda.current_device()) + ) + # Do not retain an alias while _ensure_teacher_ipc_storage drops an + # incompatible worker-owned tensor; that would keep both slabs live. + del current_storage + packed_output_storage = self._ensure_teacher_ipc_storage( + total_tokens=planned_total_stored_tokens, + vocab_size=planned_local_vocab_size, + dtype=torch.float32, + device=storage_device, + ) + elif not pack: + # Preserve the existing unpacked path. It cannot omit padding logits, + # but can still replace a known undersized slab before the forward. + current_storage = self._teacher_ipc_storage + growth_geometry = None + if ( + current_storage is not None + and current_storage.ndim == 2 + and current_storage.shape[0] < planned_total_stored_tokens + ): + growth_geometry = ( + int(current_storage.shape[1]), + current_storage.dtype, + current_storage.device, + ) + del current_storage + if growth_geometry is not None: + vocab_size, storage_dtype, storage_device = growth_geometry + self._ensure_teacher_ipc_storage( + total_tokens=planned_total_stored_tokens, + vocab_size=vocab_size, + dtype=storage_dtype, + device=storage_device, + ) + + ( + mb_iterator, + num_microbatches, + forward_mbs, + seq_length, + padded_seq_length, + ) = get_microbatch_iterator( + data, + self.cfg, + forward_batch_size, + straggler_timer=self.mcore_state.straggler_timer, + delegate_pack_to_model=self.delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=self.delegate_mtp_loss_mask_to_model, + model_slices_context_parallel_inputs=self.model_slices_context_parallel_inputs, + mtp_enabled=self.mtp_enabled, + ) + + batch_item_ids_by_microbatch: list[Optional[list[int]]] = [] + valid_lengths_by_microbatch: list[list[int]] = [] + + def recording_iterator(): + for processed_mb in mb_iterator: + logical_batch_size = processed_mb.data_dict.size + batch_item_ids_by_microbatch.append( + extract_batch_item_ids( + processed_mb.data_dict, + logical_batch_size, + required=self.cfg["sequence_packing"]["enabled"], + ) + ) + input_ids = processed_mb.data_dict.get("input_ids") + if not isinstance(input_ids, torch.Tensor) or input_ids.ndim < 2: + raise ValueError( + "Dense teacher IPC requires logical input_ids on every " + "processed microbatch." + ) + valid_lengths_by_microbatch.append( + extract_teacher_ipc_valid_lengths( + processed_mb.data_dict, + logical_batch_size, + full_seq_len=int(input_ids.shape[1]), + compact_padding=self.cfg["sequence_packing"]["enabled"], + ) + ) + yield processed_mb + + full_logits_postprocessor = FullLogitsPostProcessor( + cfg=self.cfg, + packed_output_storage=packed_output_storage, + ) + list_of_outputs = megatron_forward_backward( + model=self.model, + data_iterator=recording_iterator(), + seq_length=padded_seq_length, + mbs=forward_mbs, + num_microbatches=num_microbatches, + post_processing_fn=full_logits_postprocessor, + forward_only=True, + defer_fp32_logits=self.defer_fp32_logits, + sampling_params=None, + straggler_timer=self.mcore_state.straggler_timer, + model_slices_context_parallel_inputs=self.model_slices_context_parallel_inputs, + ) + + tp_rank = parallel_state.get_tensor_model_parallel_rank() + dp_rank = parallel_state.get_data_parallel_rank() + world_rank = torch.distributed.get_rank() + # The scheduler's ``padded_seq_length`` is the physical packed-bin width; + # it is intentionally not IPC geometry. The consumer's masks/alignment + # use the logical ``data['input_ids']`` axis. Packed post-processing writes + # valid prefixes directly into the persistent slab; unpacked processing + # continues to return logical rectangles. + + per_sample_handles: list[dict[str, Any]] = [] + storage: Optional[torch.Tensor] = None + token_partitions: list[list[tuple[int, torch.Tensor]]] = [] + local_valid_lengths_by_microbatch: list[list[int]] = [] + total_stored_tokens = 0 + row_offsets: list[int] = [] + payload_ipc: Optional[tuple[Any, ...]] = None + batch_sizes: list[int] = [] + local_vocab_size = 0 + streamed_metadata: list[StreamedFullLogitsMetadata] = [] + if stream_packed_outputs and not list_of_outputs: + raise RuntimeError( + "The packed teacher IPC output stage produced no microbatch metadata." + ) + if list_of_outputs: + if len(batch_item_ids_by_microbatch) != len(list_of_outputs) or len( + valid_lengths_by_microbatch + ) != len(list_of_outputs): + raise ValueError( + "Dense teacher IPC output count does not match the consumed " + "microbatch metadata count: " + f"outputs={len(list_of_outputs)}, " + f"identities={len(batch_item_ids_by_microbatch)}, " + f"valid_lengths={len(valid_lengths_by_microbatch)}." + ) + local_valid_lengths_by_microbatch = [ + localize_teacher_ipc_valid_lengths( + valid_lengths, + full_seq_len=full_seq_len, + cp_rank=cp_rank, + cp_size=cp_size, + ) + for valid_lengths in valid_lengths_by_microbatch + ] + total_stored_tokens = sum( + sum(valid_lengths) + for valid_lengths in local_valid_lengths_by_microbatch + ) + if total_stored_tokens != planned_total_stored_tokens: + raise ValueError( + "Dense teacher IPC planned token capacity does not match " + "the consumed microbatches: " + f"planned={planned_total_stored_tokens}, " + f"actual={total_stored_tokens}." + ) + + if stream_packed_outputs: + if packed_output_storage is None or planned_local_vocab_size is None: + raise RuntimeError( + "Packed teacher IPC streaming reached the output stage " + "without preallocated storage geometry." + ) + storage = packed_output_storage + local_vocab_size = planned_local_vocab_size + next_streamed_token_offset = 0 + for output_index, out in enumerate(list_of_outputs): + metadata = out.get("full_logits") + if not isinstance(metadata, StreamedFullLogitsMetadata): + raise TypeError( + "Packed teacher IPC must return tensor-free streamed " + f"metadata, got output={output_index}, " + f"type={type(metadata).__name__}." + ) + expected_lengths = tuple( + local_valid_lengths_by_microbatch[output_index] + ) + expected_batch_size = len(expected_lengths) + if ( + metadata.batch_size != expected_batch_size + or metadata.logical_seq_len != target_local_seq + or metadata.local_vocab_size != local_vocab_size + or metadata.dtype != storage.dtype + or metadata.device != storage.device + or metadata.storage_token_offset != next_streamed_token_offset + or metadata.stored_seq_lengths != expected_lengths + ): + raise ValueError( + "Packed teacher IPC streamed metadata disagrees with " + "the planned logical geometry: " + f"output={output_index}, metadata={metadata}, " + f"expected_batch={expected_batch_size}, " + f"expected_seq={target_local_seq}, " + f"expected_vocab={local_vocab_size}, " + f"expected_dtype={storage.dtype}, " + f"expected_device={storage.device}, " + f"expected_offset={next_streamed_token_offset}, " + f"expected_lengths={expected_lengths}." + ) + streamed_metadata.append(metadata) + batch_sizes.append(expected_batch_size) + next_streamed_token_offset += sum(expected_lengths) + if ( + next_streamed_token_offset != total_stored_tokens + or full_logits_postprocessor.packed_output_token_cursor + != total_stored_tokens + ): + raise ValueError( + "Packed teacher IPC streaming did not fill the planned slab " + "exactly: " + f"metadata_tokens={next_streamed_token_offset}, " + f"cursor={full_logits_postprocessor.packed_output_token_cursor}, " + f"planned={total_stored_tokens}." + ) + else: + first_vals = list_of_outputs[0]["full_logits"] + if not isinstance(first_vals, torch.Tensor) or first_vals.ndim != 3: + raise ValueError( + "Dense teacher IPC expects three-dimensional logits, got " + f"type={type(first_vals).__name__}, " + f"shape={getattr(first_vals, 'shape', None)}." + ) + local_vocab_size = int(first_vals.shape[-1]) + for output_index, out in enumerate(list_of_outputs): + vals = out["full_logits"] + if not isinstance(vals, torch.Tensor) or vals.ndim != 3: + raise ValueError( + "Dense teacher IPC expects three-dimensional logits, " + f"got output={output_index}, " + f"type={type(vals).__name__}, " + f"shape={getattr(vals, 'shape', None)}." + ) + batch_size_mb, seq_len_mb, output_vocab_size = vals.shape + if batch_size_mb <= 0: + raise ValueError( + "Dense teacher IPC microbatches must contain at least " + f"one logical row, got output={output_index}." + ) + if ( + seq_len_mb > target_local_seq + or output_vocab_size != local_vocab_size + or vals.dtype != first_vals.dtype + or vals.device != first_vals.device + ): + raise ValueError( + "Dense teacher IPC outputs must share sequence/vocabulary " + "geometry, dtype, and device: " + f"output={output_index}, value={tuple(vals.shape)}/" + f"{vals.dtype}/{vals.device}, expected seq_len<=" + f"{target_local_seq}, vocab={local_vocab_size}, " + f"where={first_vals.dtype}/{first_vals.device}." + ) + if len(valid_lengths_by_microbatch[output_index]) != batch_size_mb: + raise ValueError( + "Dense teacher IPC valid-length cardinality does not " + f"match output={output_index}: lengths=" + f"{len(valid_lengths_by_microbatch[output_index])}, " + f"rows={batch_size_mb}." + ) + batch_sizes.append(int(batch_size_mb)) + storage = self._ensure_teacher_ipc_storage( + total_tokens=total_stored_tokens, + vocab_size=local_vocab_size, + dtype=first_vals.dtype, + device=first_vals.device, + ) + + token_partitions = partition_teacher_ipc_token_buffer( + storage, local_valid_lengths_by_microbatch + ) + if total_stored_tokens > 0: + payload_ipc = get_handle_from_tensor(storage) + self._teacher_ipc_handles = [payload_ipc] + else: + self._teacher_ipc_handles = [] + row_cursor = 0 + for batch_size in batch_sizes: + row_offsets.append(row_cursor) + row_cursor += batch_size + + for output_index, out in enumerate(list_of_outputs): + vals = None + if stream_packed_outputs: + metadata = streamed_metadata[output_index] + batch_size_mb = metadata.batch_size + seq_len_mb = metadata.logical_seq_len + else: + # Drop each dense forward-output owner immediately after its copy. + vals = out.pop("full_logits") + assert isinstance(vals, torch.Tensor) + batch_size_mb, seq_len_mb, _ = vals.shape + assert storage is not None + row_partitions = token_partitions[output_index] + local_valid_lengths = local_valid_lengths_by_microbatch[output_index] + if len(row_partitions) != batch_size_mb: + raise ValueError( + "Dense teacher IPC output does not match its compact token " + f"partition: rows={batch_size_mb}, partitions=" + f"{len(row_partitions)}." + ) + + batch_item_ids = batch_item_ids_by_microbatch[output_index] + full_vocab_size = local_vocab_size * tp_size + vocab_start_index = tp_rank * local_vocab_size + vocab_end_index = (tp_rank + 1) * local_vocab_size + for sample_index_in_buf in range(batch_size_mb): + token_offset, row_view = row_partitions[sample_index_in_buf] + stored_seq_len = local_valid_lengths[sample_index_in_buf] + if stored_seq_len > seq_len_mb or tuple(row_view.shape) != ( + stored_seq_len, + local_vocab_size, + ): + raise ValueError( + "Dense teacher IPC valid-prefix partition exceeds its " + f"source row: output={output_index}, sample=" + f"{sample_index_in_buf}, stored={stored_seq_len}, " + f"source_shape={None if vals is None else tuple(vals.shape)}, " + "view_shape=" + f"{tuple(row_view.shape)}." + ) + if vals is not None and stored_seq_len > 0: + row_view.copy_(vals[sample_index_in_buf, :stored_seq_len]) + handle_record = { + "payload_ipc": payload_ipc, + "buf_idx": 0, + "sample_index_in_buf": sample_index_in_buf, + "ipc_layout": DENSE_TEACHER_IPC_FLAT_LAYOUT, + "storage_shape": tuple(storage.shape), + "storage_row_offset": row_offsets[output_index], + "storage_capacity_rows": sum(batch_sizes), + "storage_token_offset": token_offset, + "storage_used_tokens": total_stored_tokens, + "storage_capacity_tokens": int(storage.shape[0]), + "stored_seq_len": stored_seq_len, + "valid_seq_len": valid_lengths_by_microbatch[output_index][ + sample_index_in_buf + ], + "actual_shape": (target_local_seq, local_vocab_size), + "dtype": storage.dtype, + "tp_rank": tp_rank, + "cp_rank": cp_rank, + "tp_size": tp_size, + "cp_size": cp_size, + "world_rank": world_rank, + "vocab_start_index": vocab_start_index, + "vocab_end_index": vocab_end_index, + "global_seq_start": global_seq_start, + "full_vocab_size": full_vocab_size, + "full_seq_len": full_seq_len, + "vocab_sharded": tp_size > 1, + "sequence_sharded": cp_size > 1, + } + if batch_item_ids is not None: + handle_record["batch_item_id"] = batch_item_ids[sample_index_in_buf] + per_sample_handles.append(handle_record) + if vals is not None: + del vals + # Force the async copies above to complete before the IPC handles are + # consumed by the student process, so the consumer can't observe a + # partially written buffer. + torch.cuda.synchronize() + return {"per_sample_handles": per_sample_handles, "dp_rank": dp_rank} + + @torch.no_grad() + @wrap_with_nvtx_name("megatron_policy_worker/get_sparse_logits_ipc") + def get_sparse_logits_ipc( + self, + data: BatchedDataDict[GenerationDatumSpec], + *, + k: int, + temperature: float, + real_vocab_size: int, + noise_filter_k: int, + micro_batch_size: Optional[int] = None, + ) -> dict[str, Any]: + """Teacher forward; sparse rows of each CP shard via persistent CUDA IPC. + + Every TP rank joins the exact top-k / log Z collectives, but the result + is TP-replicated, so only TP rank 0 on the last pipeline stage exports + storage. Rows stay in their CP shard: one record is emitted per teacher + CP rank and sample, and ``sequence_segments`` lists the global + intervals its rows cover (see :func:`cp_load_balanced_segments`). + + ``data`` must carry ``force_include_token_ids`` ``[B, T_t, F]``. Each + record holds seven IPC handles (natural top-k logits and ids sorted by + id, log Z, the score-ranked k-th id, and the forced-label logits, ids + and top-k membership) into slabs that are reused across calls and freed + by :meth:`release_ipc_buffer`. + + Returns: + ``{"per_sample_handles": list, "dp_rank": int}``. + """ + if ( + self.delegate_pack_to_model + and parallel_state.get_context_parallel_world_size() > 1 + ): + raise NotImplementedError( + "get_sparse_logits_ipc does not support context parallelism for " + "models that pack and context-shard inside their own forward. " + "Set context_parallel_size=1 for this teacher." + ) + forward_batch_size = ( + micro_batch_size + if micro_batch_size is not None + else self.cfg["logprob_batch_size"] + ) + self.model.eval() + ( + mb_iterator, + num_microbatches, + forward_mbs, + _seq_length, + padded_seq_length, + ) = get_microbatch_iterator( + data, + self.cfg, + forward_batch_size, + straggler_timer=self.mcore_state.straggler_timer, + delegate_pack_to_model=self.delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=self.delegate_mtp_loss_mask_to_model, + model_slices_context_parallel_inputs=self.model_slices_context_parallel_inputs, + mtp_enabled=self.mtp_enabled, + skip_keys={"force_include_token_ids"}, + ) + list_of_outputs = megatron_forward_backward( + model=self.model, + data_iterator=mb_iterator, + seq_length=padded_seq_length, + mbs=forward_mbs, + num_microbatches=num_microbatches, + post_processing_fn=SparseTeacherLogitsPostProcessor( + cfg=self.cfg, + k=k, + temperature=temperature, + real_vocab_size=real_vocab_size, + noise_filter_k=noise_filter_k, + ), + forward_only=True, + defer_fp32_logits=self.defer_fp32_logits, + sampling_params=None, + straggler_timer=self.mcore_state.straggler_timer, + model_slices_context_parallel_inputs=self.model_slices_context_parallel_inputs, + ) + + tp_rank = parallel_state.get_tensor_model_parallel_rank() + tp_size = parallel_state.get_tensor_model_parallel_world_size() + dp_rank = parallel_state.get_data_parallel_rank() + if ( + not parallel_state.is_pipeline_last_stage(ignore_virtual=True) + or tp_rank != 0 + ): + return {"per_sample_handles": [], "dp_rank": dp_rank} + + cp_rank = parallel_state.get_context_parallel_rank() + cp_size = parallel_state.get_context_parallel_world_size() + world_rank = torch.distributed.get_rank() + target_local_seq = padded_seq_length // cp_size + # Reserve the configured maximum length on first allocation so a later, + # longer batch reuses the slabs instead of reallocating next to them. + ipc_full_seq_capacity = max( + padded_seq_length, int(self.cfg["max_total_sequence_length"]) + ) + ipc_local_seq_capacity = (ipc_full_seq_capacity + cp_size - 1) // cp_size + if not self._teacher_sparse_ipc_storage: + print( + "Sparse teacher IPC capacity: " + f"local_seq={target_local_seq} " + f"capacity_local_seq={ipc_local_seq_capacity} " + f"num_microbatches={num_microbatches} k={k}", + flush=True, + ) + sequence_segments = cp_load_balanced_segments( + padded_seq_length, cp_rank, cp_size + ) + + def _ensure_field( + name: str, + tensor: torch.Tensor, + *, + batch_size: int, + width: int, + ) -> tuple[torch.Tensor, tuple[Any, ...]]: + storage, handle = ensure_teacher_ipc_buffer( + self._teacher_sparse_ipc_storage.get(name), + self._teacher_sparse_ipc_handles.get(name), + num_microbatches, + batch_size, + target_local_seq, + width, + tensor.dtype, + tensor.device, + capacity_seq_len=ipc_local_seq_capacity, + ) + self._teacher_sparse_ipc_storage[name] = storage + self._teacher_sparse_ipc_handles[name] = handle + return storage, handle + + # Megatron pads the vocabulary uniformly across TP; the sparse result + # does not carry the local vocab width, so read it from the model. + padded_global_vocab = int(self.final_padded_vocab_size) + per_sample_handles: list[dict[str, Any]] = [] + for buf_idx, out in enumerate(list_of_outputs): + topk_logits = out["sparse_topk_logits"] + topk_indices = out["sparse_topk_indices"] + log_z = out["sparse_log_z"].unsqueeze(-1) + natural_tail_indices = out["sparse_natural_tail_indices"] + forced_logits = out["sparse_forced_logits"] + forced_indices = out["sparse_forced_indices"] + forced_in_topk = out["sparse_forced_in_topk"] + batch_size_mb, seq_len_mb, topk_width = topk_logits.shape + force_width = int(forced_logits.shape[-1]) + if ( + natural_tail_indices.shape != (batch_size_mb, seq_len_mb, 1) + or forced_indices.shape != forced_logits.shape + or forced_in_topk.shape != forced_logits.shape + or forced_logits.shape[:2] != (batch_size_mb, seq_len_mb) + or force_width <= 0 + ): + raise ValueError( + "Sparse teacher natural/forced sidecar shapes disagree: " + f"topk={tuple(topk_logits.shape)}, " + f"tail={tuple(natural_tail_indices.shape)}, " + f"forced_logits={tuple(forced_logits.shape)}, " + f"forced_indices={tuple(forced_indices.shape)}, " + f"forced_in_topk={tuple(forced_in_topk.shape)}." + ) + if seq_len_mb != target_local_seq: + raise ValueError( + f"Sparse teacher local sequence {seq_len_mb} does not match " + f"the padded CP length {target_local_seq}." + ) + + fields = { + "topk_logits": (topk_logits, topk_width), + "topk_indices": (topk_indices, topk_width), + "log_z": (log_z, 1), + "natural_tail_indices": (natural_tail_indices, 1), + "forced_logits": (forced_logits, force_width), + "forced_indices": (forced_indices, force_width), + "forced_in_topk": (forced_in_topk, force_width), + } + ipc_handles: dict[str, tuple[Any, ...]] = {} + for name, (value, width) in fields.items(): + storage, handle = _ensure_field( + name, value, batch_size=batch_size_mb, width=width + ) + storage[buf_idx, :batch_size_mb, :target_local_seq, :width].copy_(value) + ipc_handles[name] = handle + + for sample_index_in_buf in range(batch_size_mb): + per_sample_handles.append( + { + "transport": "sparse_topk", + "topk_logits_ipc": ipc_handles["topk_logits"], + "topk_indices_ipc": ipc_handles["topk_indices"], + "log_z_ipc": ipc_handles["log_z"], + "natural_tail_indices_ipc": ipc_handles["natural_tail_indices"], + "forced_logits_ipc": ipc_handles["forced_logits"], + "forced_indices_ipc": ipc_handles["forced_indices"], + "forced_in_topk_ipc": ipc_handles["forced_in_topk"], + "buf_idx": buf_idx, + "sample_index_in_buf": sample_index_in_buf, + "topk_shape": (target_local_seq, topk_width), + "forced_shape": (target_local_seq, force_width), + "scalar_shape": (target_local_seq,), + "k": int(k), + "force_width": force_width, + "temperature": float(temperature), + "real_vocab_size": int(real_vocab_size), + "padded_vocab_size": padded_global_vocab, + "forced_in_topk_k": int(noise_filter_k or k), + "tp_rank": tp_rank, + "tp_size": tp_size, + "cp_rank": cp_rank, + "cp_size": cp_size, + "world_rank": world_rank, + "full_seq_len": padded_seq_length, + "sequence_segments": sequence_segments, + } + ) + # The student imports these slabs as soon as this call returns; finish + # every copy first so it cannot observe a partially written row. + torch.cuda.synchronize() + return {"per_sample_handles": per_sample_handles, "dp_rank": dp_rank} + + def release_ipc_buffer(self) -> None: + """Free the persistent teacher-logit IPC storage. + + Called once at the end of training / validation (and on error) by the + driver, not per call. + """ + self._teacher_ipc_storage = None + self._teacher_ipc_handles = [] + self._teacher_sparse_ipc_storage = {} + self._teacher_sparse_ipc_handles = {} + gc.collect() + torch.cuda.empty_cache() + @torch.no_grad() @wrap_with_nvtx_name("megatron_policy_worker/prepare_refit_info") def _prepare_source_refit_info(self) -> dict[str, tuple[torch.Size, torch.dtype]]: @@ -4953,7 +5738,13 @@ def move_model( else: # Ordinary offload case if move_params: - model.to(device=device, non_blocking=True) + # Inference-only Megatron teachers are not wrapped in DDP. + # Loading a CPU state dict into a CUDA module only copies the + # values back into its existing CUDA parameters; it does not + # change their device and therefore does not free GPU memory. + # Move the module itself, matching the Automodel worker's + # offload/onload lifecycle. + model = model.to(device=device, non_blocking=True) return model def move_optimizer(self, device: str): diff --git a/nemo_rl/models/value/workers/automodel_value_worker.py b/nemo_rl/models/value/workers/automodel_value_worker.py index 2044b8f603e..75c6c195645 100644 --- a/nemo_rl/models/value/workers/automodel_value_worker.py +++ b/nemo_rl/models/value/workers/automodel_value_worker.py @@ -246,6 +246,7 @@ def __init__( self.is_generation_colocated, _runtime_sampling_params, _runtime_is_reward_model, + _runtime_model_load_dtype, ) = runtime_config def _autocast_context(self) -> AbstractContextManager[Any]: diff --git a/pyrefly.toml b/pyrefly.toml index d21f60cc856..899d9ae7e61 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -25,6 +25,11 @@ replace-imports-with-any = [ "zstandard.*", ] project-includes = [ + "nemo_rl/algorithms/loss/csd.py", + "nemo_rl/algorithms/loss/csd_packing.py", + "tools/csd/__init__.py", + "tools/csd/build_subtoks_table.py", + "tools/csd/validate_subtoks_tables.py", "nemo_rl/data/captured_media.py", "nemo_rl/experience/reward_penalties.py", # TODO: enable these once we have 100 correctness @@ -88,6 +93,7 @@ project-includes = [ "nemo_rl/data/cross_tokenizer/__init__.py", "nemo_rl/data/cross_tokenizer/native_chat.py", "nemo_rl/data/cross_tokenizer/cross_tokenizer_collate.py", + "nemo_rl/data/cross_tokenizer/csd_collate.py", "nemo_rl/data/collate_fn.py", "nemo_rl/data/dataloader.py", "nemo_rl/data/datasets/__init__.py", @@ -153,6 +159,7 @@ project-includes = [ "nemo_rl/data/multimodal_utils.py", "nemo_rl/data/megatron_sft_packed.py", "nemo_rl/data/packing/__init__.py", + "nemo_rl/data/packing/lockstep.py", "nemo_rl/data/processors.py", "nemo_rl/data_plane/__init__.py", "nemo_rl/data_plane/adapters/__init__.py", @@ -226,6 +233,7 @@ project-includes = [ "nemo_rl/models/automodel/convert.py", "nemo_rl/models/automodel/router_replay.py", "nemo_rl/models/automodel/utils.py", + "nemo_rl/models/automodel/nemotron_h_cp.py", "nemo_rl/models/generation/__init__.py", "nemo_rl/models/generation/constants.py", "nemo_rl/models/generation/dynamo/__init__.py", diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index 58bad617fb6..be4616b22f5 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -94,3 +94,7 @@ research/template_project/tests/test_suites/llm/single_update_1n8g.sh # Full Flow-GRPO recipe requires an externally served GenRM OCR judge. research/flow_grpo/tests/test_suites/flow_grpo_qwen_image_ocr.sh + +# CSD single-teacher run on the Automodel backend. The Megatron CSD nightly +# covers the default CSD example; run this one manually to check Automodel. +tests/test_suites/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.sh diff --git a/tests/test_suites/llm/csd-multiteacher-qwen3-4b-llama3.2-3b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.sh b/tests/test_suites/llm/csd-multiteacher-qwen3-4b-llama3.2-3b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.sh new file mode 100755 index 00000000000..e18b8fbbdd1 --- /dev/null +++ b/tests/test_suites/llm/csd-multiteacher-qwen3-4b-llama3.2-3b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.sh @@ -0,0 +1,51 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=1 +STEPS_PER_RUN=100 +MAX_STEPS=100 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=30 +STUDENT_MODEL=meta-llama/Llama-3.2-1B +TEACHER0_MODEL=Qwen/Qwen3-4B # cross-tokenizer (P-KL via projection) +TEACHER1_MODEL=meta-llama/Llama-3.2-3B # same tokenizer/vocab as student +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# Prepare both complete token-chain directions for the cross-tokenizer teacher. +TABLE_DIR="$EXP_DIR/csd_tables" +mkdir -p "$TABLE_DIR" +CSD_STUDENT_TO_TEACHER_TABLE="$TABLE_DIR/student_to_teacher.pt" +CSD_TEACHER_TO_STUDENT_TABLE="$TABLE_DIR/teacher_to_student.pt" +uv run python -m tools.csd.build_subtoks_table \ + --source-tokenizer "$STUDENT_MODEL" --destination-tokenizer "$TEACHER0_MODEL" \ + --output "$CSD_STUDENT_TO_TEACHER_TABLE" --overwrite +uv run python -m tools.csd.build_subtoks_table \ + --source-tokenizer "$TEACHER0_MODEL" --destination-tokenizer "$STUDENT_MODEL" \ + --output "$CSD_TEACHER_TO_STUDENT_TABLE" --overwrite + +uv run examples/run_csd.py \ + --config $CONFIG_PATH \ + distillation.max_num_steps=$MAX_STEPS \ + teachers.0.aligner.student_to_teacher_table_path="$CSD_STUDENT_TO_TEACHER_TABLE" \ + teachers.0.aligner.teacher_to_student_table_path="$CSD_TEACHER_TO_STUDENT_TABLE" \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + $@ \ + 2>&1 | tee $RUN_LOG + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Smoke qualification checks finiteness; it does not assert benchmark convergence. +uv run tests/check_metrics.py "$JSON_METRICS" 'all_finite(data["train/loss"])' diff --git a/tests/test_suites/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.sh b/tests/test_suites/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.sh new file mode 100755 index 00000000000..db3810a89e3 --- /dev/null +++ b/tests/test_suites/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.sh @@ -0,0 +1,50 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=1 +STEPS_PER_RUN=100 +MAX_STEPS=100 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=20 +STUDENT_MODEL=meta-llama/Llama-3.2-1B +TEACHER_MODEL=Qwen/Qwen3-4B +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# Prepare both complete token-chain directions for the cross-tokenizer teacher. +TABLE_DIR="$EXP_DIR/csd_tables" +mkdir -p "$TABLE_DIR" +CSD_STUDENT_TO_TEACHER_TABLE="$TABLE_DIR/student_to_teacher.pt" +CSD_TEACHER_TO_STUDENT_TABLE="$TABLE_DIR/teacher_to_student.pt" +uv run python -m tools.csd.build_subtoks_table \ + --source-tokenizer "$STUDENT_MODEL" --destination-tokenizer "$TEACHER_MODEL" \ + --output "$CSD_STUDENT_TO_TEACHER_TABLE" --overwrite +uv run python -m tools.csd.build_subtoks_table \ + --source-tokenizer "$TEACHER_MODEL" --destination-tokenizer "$STUDENT_MODEL" \ + --output "$CSD_TEACHER_TO_STUDENT_TABLE" --overwrite + +uv run examples/run_csd.py \ + --config $CONFIG_PATH \ + distillation.max_num_steps=$MAX_STEPS \ + teachers.0.aligner.student_to_teacher_table_path="$CSD_STUDENT_TO_TEACHER_TABLE" \ + teachers.0.aligner.teacher_to_student_table_path="$CSD_TEACHER_TO_STUDENT_TABLE" \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + $@ \ + 2>&1 | tee $RUN_LOG + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Smoke qualification checks finiteness; it does not assert benchmark convergence. +uv run tests/check_metrics.py "$JSON_METRICS" 'all_finite(data["train/loss"])' diff --git a/tests/test_suites/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-megatron-tp2cp2.sh b/tests/test_suites/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-megatron-tp2cp2.sh new file mode 100755 index 00000000000..db3810a89e3 --- /dev/null +++ b/tests/test_suites/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-megatron-tp2cp2.sh @@ -0,0 +1,50 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=1 +STEPS_PER_RUN=100 +MAX_STEPS=100 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +NUM_MINUTES=20 +STUDENT_MODEL=meta-llama/Llama-3.2-1B +TEACHER_MODEL=Qwen/Qwen3-4B +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# Prepare both complete token-chain directions for the cross-tokenizer teacher. +TABLE_DIR="$EXP_DIR/csd_tables" +mkdir -p "$TABLE_DIR" +CSD_STUDENT_TO_TEACHER_TABLE="$TABLE_DIR/student_to_teacher.pt" +CSD_TEACHER_TO_STUDENT_TABLE="$TABLE_DIR/teacher_to_student.pt" +uv run python -m tools.csd.build_subtoks_table \ + --source-tokenizer "$STUDENT_MODEL" --destination-tokenizer "$TEACHER_MODEL" \ + --output "$CSD_STUDENT_TO_TEACHER_TABLE" --overwrite +uv run python -m tools.csd.build_subtoks_table \ + --source-tokenizer "$TEACHER_MODEL" --destination-tokenizer "$STUDENT_MODEL" \ + --output "$CSD_TEACHER_TO_STUDENT_TABLE" --overwrite + +uv run examples/run_csd.py \ + --config $CONFIG_PATH \ + distillation.max_num_steps=$MAX_STEPS \ + teachers.0.aligner.student_to_teacher_table_path="$CSD_STUDENT_TO_TEACHER_TABLE" \ + teachers.0.aligner.teacher_to_student_table_path="$CSD_TEACHER_TO_STUDENT_TABLE" \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + $@ \ + 2>&1 | tee $RUN_LOG + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Smoke qualification checks finiteness; it does not assert benchmark convergence. +uv run tests/check_metrics.py "$JSON_METRICS" 'all_finite(data["train/loss"])' diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 5f609b66e9c..0f10ef9d853 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -321,6 +321,12 @@ tests/test_suites/llm/distillation-xtoken-off-policy-qwen3-4b-to-llama3.2-1b-1n8 # Multi-teacher cross-tokenizer off-policy distillation (xtoken): Qwen3-4B (cross-tok P-KL) + Llama-3.2-3B (same-vocab direct KL) -> Llama-3.2-1B, kd_loss_mode=sum, student TP4xCP2 <- teachers TP2xCP2, guards multi-teacher aggregation + per-teacher CP-buffer sharding. tests/test_suites/llm/distillation-xtoken-off-policy-multiteacher-qwen3-4b-llama3.2-3b-to-llama3.2-1b-1n8g-automodel-tp4cp2.sh +# CSD on Megatron: Qwen3-4B -> Llama-3.2-1B on raw text, student and teacher TP2xCP2, guards finite training. +tests/test_suites/llm/csd-qwen3-4b-to-llama3.2-1b-1n8g-megatron-tp2cp2.sh + +# Multi-teacher CSD: Qwen3-4B (cross-tokenizer CSD) + Llama-3.2-3B (same-vocab direct KL) -> Llama-3.2-1B, kd_loss_mode=sum, student TP4xCP2 <- teachers TP2xCP2, guards finite training. +tests/test_suites/llm/csd-multiteacher-qwen3-4b-llama3.2-3b-to-llama3.2-1b-1n8g-dtensor-tp4cp2.sh + ####### # PPO # ####### diff --git a/tests/unit/algorithms/csd/__init__.py b/tests/unit/algorithms/csd/__init__.py new file mode 100644 index 00000000000..4f17b0780d0 --- /dev/null +++ b/tests/unit/algorithms/csd/__init__.py @@ -0,0 +1 @@ +"""CSD controller, loss, and distributed regression tests.""" diff --git a/tests/unit/algorithms/csd/test_ce_sharded_student_memory.py b/tests/unit/algorithms/csd/test_ce_sharded_student_memory.py new file mode 100644 index 00000000000..eb9388cf5bd --- /dev/null +++ b/tests/unit/algorithms/csd/test_ce_sharded_student_memory.py @@ -0,0 +1,239 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import tempfile +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + + +def _run_tp2_cp2_ce_memory( + rank: int, + world_size: int, + init_file: str, +): + from nemo_rl.algorithms.loss import csd as loss_functions_module + from nemo_rl.algorithms.loss.csd import ( + CSD_STUDENT_LOSS_CHUNK_SIZE, + CSDLossFn, + ) + from nemo_rl.algorithms.x_token import loss_utils as shared_loss_utils + from nemo_rl.distributed import model_utils + from nemo_rl.distributed.batched_data_dict import BatchedDataDict + from nemo_rl.distributed.model_utils import _get_tokens_on_this_cp_rank + + torch.set_num_threads(1) + dist.init_process_group( + backend="gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + ) + try: + tp_size = cp_size = 2 + tp_index, cp_index = rank // cp_size, rank % cp_size + tp_group = None + for cp_coordinate in range(cp_size): + group = dist.new_group( + ranks=[ + tp_coordinate * cp_size + cp_coordinate + for tp_coordinate in range(tp_size) + ] + ) + if cp_coordinate == cp_index: + tp_group = group + cp_group = None + for tp_coordinate in range(tp_size): + group = dist.new_group( + ranks=[ + tp_coordinate * cp_size + cp_coordinate + for cp_coordinate in range(cp_size) + ] + ) + if tp_coordinate == tp_index: + cp_group = group + assert tp_group is not None + assert cp_group is not None + + torch.manual_seed(11) + batch_size = 2 + sequence_length = 1028 + vocab_size = 12 + full_logits = torch.randn(batch_size, sequence_length, vocab_size) + input_ids = torch.randint( + 0, + vocab_size, + (batch_size, sequence_length), + dtype=torch.long, + ) + token_mask = torch.ones(batch_size, sequence_length) + token_mask[0, 257:269] = 0 + token_mask[0, -1] = 0 + sample_mask = torch.tensor([1.0, 0.0]) + data = BatchedDataDict( + { + "input_ids": input_ids, + "token_mask": token_mask, + "sample_mask": sample_mask, + } + ) + global_valid_toks = (token_mask[:, 1:] * sample_mask.unsqueeze(-1)).sum() + loss_fn = CSDLossFn.__new__(CSDLossFn) + + reference_logits = full_logits.clone().requires_grad_(True) + reference_loss = loss_fn._compute_ce( + reference_logits, + data, + global_valid_toks, + tp_group=None, + cp_group=None, + ) + reference_loss.backward() + reference_grad = reference_logits.grad.detach().clone() + + vocab_block = vocab_size // tp_size + vocab_start = tp_index * vocab_block + vocab_slice = slice(vocab_start, vocab_start + vocab_block) + student_shard = _get_tokens_on_this_cp_rank( + full_logits[:, :, vocab_slice], + cp_rank=cp_index, + cp_size=cp_size, + seq_dim=1, + ).contiguous() + student_shard = student_shard.clone().requires_grad_(True) + assert student_shard.shape[1] > CSD_STUDENT_LOSS_CHUNK_SIZE + + original_cp_contiguous_gather = ( + loss_functions_module.allgather_cp_contiguous_tensor + ) + original_full_softmax = shared_loss_utils.vocab_parallel_full_log_softmax + original_scalar_cp_gather = model_utils.allgather_cp_sharded_tensor + had_vocab_gather = hasattr( + loss_functions_module, + "vocab_parallel_gather_logits", + ) + original_vocab_gather = getattr( + loss_functions_module, + "vocab_parallel_gather_logits", + None, + ) + scalar_gathers = 0 + + def reject_contiguous_rank3_gather(*_args, **_kwargs): + raise AssertionError("CE attempted a rank-3 contiguous CP gather") + + def reject_vocab_gather(*_args, **_kwargs): + raise AssertionError("CE attempted a rank-3 TP vocabulary gather") + + def require_scalar_cp_gather(tensor, group, seq_dim=1): + nonlocal scalar_gathers + assert tensor.ndim == 2 + scalar_gathers += 1 + return original_scalar_cp_gather(tensor, group, seq_dim) + + loss_functions_module.allgather_cp_contiguous_tensor = ( + reject_contiguous_rank3_gather + ) + shared_loss_utils.vocab_parallel_full_log_softmax = reject_vocab_gather + loss_functions_module.vocab_parallel_gather_logits = reject_vocab_gather + model_utils.allgather_cp_sharded_tensor = require_scalar_cp_gather + try: + sharded_loss = loss_fn._compute_ce( + student_shard, + data, + global_valid_toks, + student_logits_contig=student_shard, + tp_group=tp_group, + cp_group=cp_group, + ) + finally: + loss_functions_module.allgather_cp_contiguous_tensor = ( + original_cp_contiguous_gather + ) + shared_loss_utils.vocab_parallel_full_log_softmax = original_full_softmax + model_utils.allgather_cp_sharded_tensor = original_scalar_cp_gather + if had_vocab_gather: + loss_functions_module.vocab_parallel_gather_logits = ( + original_vocab_gather + ) + else: + delattr(loss_functions_module, "vocab_parallel_gather_logits") + + assert scalar_gathers == 1 + torch.testing.assert_close( + sharded_loss, + reference_loss, + rtol=1.0e-5, + atol=1.0e-6, + ) + (sharded_loss / float(cp_size)).backward() + + gathered_gradients = [torch.empty_like(student_shard.grad) for _ in range(4)] + dist.all_gather(gathered_gradients, student_shard.grad.contiguous()) + reconstructed_grad = torch.zeros_like(reference_grad) + sequence_chunk = sequence_length // (2 * cp_size) + for shard_rank, shard_gradient in enumerate(gathered_gradients): + shard_tp, shard_cp = shard_rank // cp_size, shard_rank % cp_size + shard_vocab = slice( + shard_tp * vocab_block, + (shard_tp + 1) * vocab_block, + ) + head, tail = torch.chunk(shard_gradient, chunks=2, dim=1) + head_chunk = shard_cp + tail_chunk = 2 * cp_size - shard_cp - 1 + reconstructed_grad[ + :, + head_chunk * sequence_chunk : (head_chunk + 1) * sequence_chunk, + shard_vocab, + ] = head + reconstructed_grad[ + :, + tail_chunk * sequence_chunk : (tail_chunk + 1) * sequence_chunk, + shard_vocab, + ] = tail + torch.testing.assert_close( + reconstructed_grad, + reference_grad, + rtol=1.0e-5, + atol=1.0e-6, + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="gloo backend unavailable", +) +def test_ce_tp2_cp2_matches_full_loss_and_gradient_without_rank3_gathers(tmp_path): + """Masked CE is exact across TP2/CP2 and crosses the fixed chunk boundary.""" + init_file = tmp_path / "gloo_init" + mp.spawn( + _run_tp2_cp2_ce_memory, + args=(4, os.fspath(init_file)), + nprocs=4, + join=True, + ) + + +if __name__ == "__main__": + with tempfile.TemporaryDirectory() as tmpdir: + test_ce_tp2_cp2_matches_full_loss_and_gradient_without_rank3_gathers( + Path(tmpdir) + ) + print("CE TP2/CP2 MEMORY-BOUNDED PARITY PASSED") diff --git a/tests/unit/algorithms/csd/test_config.py b/tests/unit/algorithms/csd/test_config.py new file mode 100644 index 00000000000..90acd77a279 --- /dev/null +++ b/tests/unit/algorithms/csd/test_config.py @@ -0,0 +1,56 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Validate CSD artifact choices before training allocates workers.""" + +from pathlib import Path + +import pytest +from omegaconf import OmegaConf +from pydantic import ValidationError + +from nemo_rl.algorithms.xtoken_off_policy_distillation import CSDMasterConfig +from nemo_rl.utils.config import load_config, register_omegaconf_resolvers + + +@pytest.mark.parametrize( + ("common_from_tables", "projection", "cross_tokenizer", "valid"), + [ + (True, None, True, True), + (False, "projection.pt", True, True), + (False, None, True, False), + (False, None, False, True), + ], +) +def test_csd_common_event_mapping_requires_the_selected_artifact( + common_from_tables: bool, + projection: str | None, + cross_tokenizer: bool, + valid: bool, +) -> None: + register_omegaconf_resolvers() + config_path = Path(__file__).resolve().parents[4] / "examples/configs/csd.yaml" + config = OmegaConf.to_container(load_config(str(config_path)), resolve=True) + config["loss_fn"]["common_indices_from_subtoks"] = common_from_tables + aligner = config["teachers"][0]["aligner"] + aligner["projection_matrix_path"] = projection + if not cross_tokenizer: + aligner["student_to_teacher_table_path"] = None + aligner["teacher_to_student_table_path"] = None + + if valid: + validated = CSDMasterConfig.model_validate(config) + assert validated.loss_fn.common_indices_from_subtoks == common_from_tables + else: + with pytest.raises(ValidationError, match="projection_matrix_path is required"): + CSDMasterConfig.model_validate(config) diff --git a/tests/unit/algorithms/csd/test_csd_loss_functions.py b/tests/unit/algorithms/csd/test_csd_loss_functions.py new file mode 100644 index 00000000000..3e0415cb808 --- /dev/null +++ b/tests/unit/algorithms/csd/test_csd_loss_functions.py @@ -0,0 +1,773 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from nemo_rl.algorithms.loss import prepare_loss_input +from nemo_rl.algorithms.loss.csd import CSDLossFn +from nemo_rl.algorithms.x_token.loss_utils import ( + LocalizedAlignment, + same_tokenizer_topk_kl, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.model_utils import ( + _get_tokens_on_this_cp_rank, + cp_load_balanced_to_contiguous, +) + +_CT_V_STUDENT = 6 +_CT_V_TEACHER = 5 +_CT_TOPK = 2 +_CT_TEMPERATURE = 2.0 + + +def _write_ct_projection(tmp_path): + """Dense top-k projection with two strict exact maps (s0->t0, s5->t4). + + Strict exact-map rule: slot-0 weight 1.0 (any later slot has weight 0). + The four middle rows are fuzzy (0.7/0.3, both slots real) so they fall in + the uncommon partition. Result: common student {0, 5} / teacher {0, 4}; + uncommon teacher {1, 2, 3}. + """ + v_s, v_t, k = _CT_V_STUDENT, _CT_V_TEACHER, _CT_TOPK + indices = torch.full((v_s, k), -1, dtype=torch.long) + likelihoods = torch.zeros((v_s, k), dtype=torch.float32) + indices[0, 0], likelihoods[0, 0] = 0, 1.0 + indices[v_s - 1, 0], likelihoods[v_s - 1, 0] = v_t - 1, 1.0 + for s, (t1, t2) in {1: (1, 2), 2: (2, 3), 3: (3, 1), 4: (1, 3)}.items(): + indices[s, 0], indices[s, 1] = t1, t2 + likelihoods[s, 0], likelihoods[s, 1] = 0.7, 0.3 + path = tmp_path / "ct_projection.pt" + torch.save({"indices": indices, "likelihoods": likelihoods}, path) + return str(path) + + +def _ct_loss_cfg(projection_path): + # Per-teacher metadata as ``setup`` injects it: parallel lists, one entry + # per teacher (here a single teacher). + return { + "temperature": _CT_TEMPERATURE, + "vocab_topk": 4, + "reverse_kl": False, + "kl_loss_weight": 1.0, + "ce_loss_scale": 1.0, + "dynamic_loss_scaling": False, + "kd_loss_mode": "sum", + "alpha": 1.0, + "normalize_teacher_by_vocab": False, + "student_vocab_size": _CT_V_STUDENT, + "projection_matrix_paths": [projection_path], + "teacher_vocab_sizes": [_CT_V_TEACHER], + "teacher_weights": [1.0], + } + + +def test_csd_mismatched_per_teacher_lists_raises(tmp_path): + """Unequal-length per-teacher lists fail loudly at construction.""" + cfg = _ct_loss_cfg(_write_ct_projection(tmp_path)) + # Two weights but a single entry in every other per-teacher list. + cfg["teacher_weights"] = [1.0, 1.0] + with pytest.raises(ValueError, match="per-teacher lists must be equal length"): + CSDLossFn(cfg) + + +def test_normalize_teacher_by_vocab_rejected_outside_sum_mode(tmp_path): + """normalize_teacher_by_vocab is a no-op outside sum mode, so reject it there.""" + cfg = _ct_loss_cfg(_write_ct_projection(tmp_path)) + cfg["kd_loss_mode"] = "averaged_logits" + cfg["normalize_teacher_by_vocab"] = True + with pytest.raises(ValueError, match="normalize_teacher_by_vocab"): + CSDLossFn(cfg) + + +def test_csd_ce_uniform_logits_equals_log_vocab(tmp_path): + """_compute_ce on uniform logits equals log(V_student) per valid token.""" + loss_fn = CSDLossFn(_ct_loss_cfg(_write_ct_projection(tmp_path))) + b, t_s = 1, 5 + logits = torch.zeros(b, t_s, _CT_V_STUDENT) # uniform -> CE == log(V) + data = BatchedDataDict( + { + "input_ids": torch.randint(0, _CT_V_STUDENT, (b, t_s)), + "token_mask": torch.ones(b, t_s), + "sample_mask": torch.ones(b), + } + ) + gvt = (data["token_mask"][:, 1:] * data["sample_mask"].unsqueeze(-1)).sum() + ce = loss_fn._compute_ce(logits, data, gvt) + assert torch.allclose(ce, torch.log(torch.tensor(float(_CT_V_STUDENT))), atol=1e-5) + + +def test_csd_ce_megatron_uses_chunked_raw_shards(tmp_path, monkeypatch): + """Plain Megatron CE routes raw TP/CP shards without rank-3 gathers.""" + loss_fn = CSDLossFn(_ct_loss_cfg(_write_ct_projection(tmp_path))) + local_logits = torch.randn(1, 2, 3) + tp_group = object() + cp_group = object() + calls = {} + + def reject_rank3_gather(*_args, **_kwargs): + raise AssertionError("CE attempted a rank-3 CP gather") + + def fake_get_rank(group): + assert group is tp_group + return 1 + + def fake_get_world_size(group): + assert group in (tp_group, cp_group) + return 2 + + def fake_logprobs(**kwargs): + calls.update(kwargs) + return torch.full((1, 3), -torch.log(torch.tensor(float(_CT_V_STUDENT)))) + + monkeypatch.setattr( + "nemo_rl.algorithms.loss.csd.allgather_cp_contiguous_tensor", + reject_rank3_gather, + ) + monkeypatch.setattr( + "nemo_rl.algorithms.loss.csd.get_next_token_logprobs_from_logits", + fake_logprobs, + ) + monkeypatch.setattr(torch.distributed, "get_rank", fake_get_rank) + monkeypatch.setattr(torch.distributed, "get_world_size", fake_get_world_size) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + + data = BatchedDataDict( + { + "input_ids": torch.randint(0, _CT_V_STUDENT, (1, 4)), + "token_mask": torch.ones(1, 4), + "sample_mask": torch.ones(1), + } + ) + global_valid_toks = data["token_mask"][:, 1:].sum() + ce = loss_fn._compute_ce( + local_logits, + data, + global_valid_toks, + student_logits_contig=local_logits, + tp_group=tp_group, + cp_group=cp_group, + ) + + assert calls["next_token_logits"] is local_logits + assert calls["input_ids"] is data["input_ids"] + assert calls["vocab_parallel_rank"] == 1 + assert calls["vocab_parallel_group"] is tp_group + assert calls["context_parallel_group"] is cp_group + assert calls["chunk_size"] == 256 + assert torch.allclose(ce, torch.log(torch.tensor(float(_CT_V_STUDENT))), atol=1e-5) + + +def test_csd_ce_respects_sample_mask(tmp_path): + """A masked sample must not contribute to _compute_ce: the B=2 result with + sample 1 masked equals the CE over sample 0 alone.""" + loss_fn = CSDLossFn(_ct_loss_cfg(_write_ct_projection(tmp_path))) + b, t_s = 2, 5 + torch.manual_seed(0) + logits = torch.randn(b, t_s, _CT_V_STUDENT) + input_ids = torch.randint(0, _CT_V_STUDENT, (b, t_s)) + token_mask = torch.ones(b, t_s) + + data_masked = BatchedDataDict( + { + "input_ids": input_ids, + "token_mask": token_mask, + "sample_mask": torch.tensor([1.0, 0.0]), + } + ) + gvt_masked = (token_mask[:, 1:] * data_masked["sample_mask"].unsqueeze(-1)).sum() + ce_masked = loss_fn._compute_ce(logits, data_masked, gvt_masked) + + data_single = BatchedDataDict( + { + "input_ids": input_ids[:1], + "token_mask": token_mask[:1], + "sample_mask": torch.tensor([1.0]), + } + ) + gvt_single = (token_mask[:1, 1:] * data_single["sample_mask"].unsqueeze(-1)).sum() + ce_single = loss_fn._compute_ce(logits[:1], data_single, gvt_single) + + assert torch.allclose(ce_masked, ce_single, atol=1e-6) + + +def test_csd_prepare_loss_input_partitions_canonical_ce(tmp_path, monkeypatch): + """Automodel CP assigns one disjoint canonical CE window to each rank.""" + loss_fn = CSDLossFn(_ct_loss_cfg(_write_ct_projection(tmp_path))) + full_logprobs = torch.tensor([[-1.0, -2.0, -3.0]], requires_grad=True) + token_mask = torch.tensor([[1.0, 1.0, 0.0, 1.0]]) + cp_group = object() + data = BatchedDataDict( + { + "input_ids": torch.tensor([[0, 1, 2, 3]]), + "token_mask": token_mask, + "sample_mask": torch.ones(1), + } + ) + + monkeypatch.setattr( + "nemo_rl.algorithms.loss.loss_input.prepare_csd_cross_tokenizer_loss_input", + lambda *args, **kwargs: (torch.empty(0), {}, {}, {}, {}, None, cp_group, None), + ) + monkeypatch.setattr( + "nemo_rl.algorithms.loss.loss_input.get_cp_sharded_next_token_logprobs", + lambda *args, **kwargs: full_logprobs, + ) + monkeypatch.setattr("torch.distributed.get_world_size", lambda group: 2) + monkeypatch.setattr("torch.distributed.get_rank", lambda group: 1) + + loss_input, prepared_data = prepare_loss_input( + torch.empty(0), + data, + loss_fn, + context_parallel_group=cp_group, + cp_sharder=object(), + ) + + torch.testing.assert_close( + loss_input["student_next_token_logprobs"], torch.tensor([[-3.0, 0.0]]) + ) + torch.testing.assert_close( + loss_input["student_next_token_mask"], torch.tensor([[1.0, 0.0]]) + ) + assert prepared_data is data + + +def test_csd_prepare_loss_input_rejects_nondivisible_cp_window(tmp_path, monkeypatch): + """Automodel CP rejects canonical CE windows that cannot partition evenly.""" + loss_fn = CSDLossFn(_ct_loss_cfg(_write_ct_projection(tmp_path))) + next_token_logprobs = torch.tensor( + [[-1.0, -2.0, -3.0, -4.0, -5.0]], requires_grad=True + ) + cp_group = object() + data = BatchedDataDict( + { + "input_ids": torch.zeros((1, 6), dtype=torch.long), + "token_mask": torch.ones((1, 6)), + "sample_mask": torch.ones(1), + } + ) + + monkeypatch.setattr( + "nemo_rl.algorithms.loss.loss_input.prepare_csd_cross_tokenizer_loss_input", + lambda *args, **kwargs: (torch.empty(0), {}, {}, {}, {}, None, cp_group, None), + ) + monkeypatch.setattr( + "nemo_rl.algorithms.loss.loss_input.get_cp_sharded_next_token_logprobs", + lambda *args, **kwargs: next_token_logprobs, + ) + monkeypatch.setattr("torch.distributed.get_world_size", lambda group: 4) + + with pytest.raises(ValueError, match=r"sequence_length=6, cp_size=4"): + prepare_loss_input( + torch.empty(0), + data, + loss_fn, + context_parallel_group=cp_group, + cp_sharder=object(), + ) + + +def test_csd_precomputed_ce_reduces_partitioned_cp_window(tmp_path, monkeypatch): + """Precomputed CE sums CP-window values while preserving local gradients.""" + loss_fn = CSDLossFn(_ct_loss_cfg(_write_ct_projection(tmp_path))) + data = BatchedDataDict({"sample_mask": torch.ones(1)}) + next_token_logprobs = torch.tensor([[-1.0, -2.0]], requires_grad=True) + next_token_mask = torch.ones_like(next_token_logprobs) + cp_group = object() + reduce_calls = [] + + def reduce_sum(value, group): + reduce_calls.append((value.detach().clone(), group)) + return value + + monkeypatch.setattr( + "nemo_rl.algorithms.loss.csd.group_all_reduce_sum_with_grad", + reduce_sum, + ) + + ce = loss_fn._compute_ce( + torch.empty(0), + data, + torch.tensor(2.0), + student_next_token_logprobs=next_token_logprobs, + student_next_token_mask=next_token_mask, + cp_group=cp_group, + ) + + torch.testing.assert_close(ce, torch.tensor(1.5)) + ce.backward() + torch.testing.assert_close( + next_token_logprobs.grad, + torch.tensor([[-0.5, -0.5]]), + ) + assert len(reduce_calls) == 1 + torch.testing.assert_close(reduce_calls[0][0], torch.tensor(1.5)) + assert reduce_calls[0][1] is cp_group + + +def test_csd_partitioned_cp_ce_matches_cp1_value_and_gradient(tmp_path): + """Summed CP2 CE windows are numerically identical to the CP1 full CE.""" + loss_fn = CSDLossFn(_ct_loss_cfg(_write_ct_projection(tmp_path))) + data = BatchedDataDict( + { + "token_mask": torch.tensor([[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0]]), + "sample_mask": torch.tensor([1.0, 0.0]), + } + ) + global_valid_toks = torch.tensor(2.0) + full_mask = data["token_mask"].roll(shifts=-1, dims=1) + full_mask[:, -1] = 0 + + cp1_logprobs = torch.tensor( + [[-1.0, -2.0, -3.0], [-10.0, -20.0, -30.0]], requires_grad=True + ) + cp1_loss = loss_fn._compute_ce( + torch.empty(0), + data, + global_valid_toks, + student_next_token_logprobs=cp1_logprobs, + student_next_token_mask=full_mask[:, :-1], + cp_group=None, + ) + cp1_loss.backward() + + cp2_logprobs = cp1_logprobs.detach().clone().requires_grad_(True) + cp2_padded = torch.cat([cp2_logprobs, torch.zeros_like(cp2_logprobs[:, :1])], dim=1) + cp2_local_seq_len = cp2_padded.shape[1] // 2 + cp2_losses = [] + for cp_rank in range(2): + seq_start = cp_rank * cp2_local_seq_len + cp2_losses.append( + loss_fn._compute_ce( + torch.empty(0), + data, + global_valid_toks, + student_next_token_logprobs=cp2_padded.narrow( + 1, seq_start, cp2_local_seq_len + ), + student_next_token_mask=full_mask.narrow( + 1, seq_start, cp2_local_seq_len + ), + cp_group=None, + ) + ) + cp2_loss = sum(cp2_losses) + cp2_loss.backward() + + torch.testing.assert_close(cp2_loss, cp1_loss) + torch.testing.assert_close(cp2_logprobs.grad, cp1_logprobs.grad) + + +def test_csd_same_vocab_topk_ignores_masked_and_padding_rows(tmp_path): + """Invalid teacher rows below the valid column maxima cannot change same-vocab loss or gradient.""" + cfg = _ct_loss_cfg(str(tmp_path / "unused_projection.pt")) + cfg["projection_matrix_paths"] = [None] + cfg["teacher_vocab_sizes"] = [_CT_V_STUDENT] + cfg["vocab_topk"] = 2 + cfg["temperature"] = 1.0 + cfg["ce_loss_scale"] = 0.0 + loss_fn = CSDLossFn(cfg) + + torch.manual_seed(7) + student_base = torch.randn(2, 4, _CT_V_STUDENT) + input_ids = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 4]]) + token_mask = torch.tensor([[1, 1, 1, 0], [1, 1, 1, 1]]) + sample_mask = torch.tensor([True, False]) + data = BatchedDataDict( + { + "input_ids": input_ids, + "input_lengths": torch.tensor([3, 4]), + "token_mask": token_mask, + "sample_mask": sample_mask, + } + ) + align = LocalizedAlignment( + sample_mask=sample_mask, + student_input_ids=input_ids, + student_token_mask=token_mask, + ) + + teacher_clean = torch.full((2, 4, _CT_V_STUDENT), -10.0) + teacher_clean[0, 0, 0:2] = torch.tensor([8.0, 7.0]) + teacher_clean[0, 1, 0:2] = torch.tensor([7.0, 8.0]) + teacher_corrupt = teacher_clean.clone() + # The top-K columns come from the per-vocab max over every teacher row, so + # the corruption stays below the valid rows' column maxima (8.0) and the + # selected columns {0, 1} are unchanged. Padding in the valid sample and + # every predictor in the masked sample still get a different distribution + # over those columns and over the rest of the vocab. + teacher_corrupt[0, 2:, 0:4] = torch.tensor([7.5, -5.0, 6.0, 5.0]) + teacher_corrupt[1, :, 0:6] = torch.tensor([-5.0, 7.5, -8.0, -9.0, 6.0, 5.0]) + + def run(teacher_logits): + student_logits = student_base.clone().requires_grad_(True) + loss, metrics = loss_fn( + data=data, + global_valid_seqs=sample_mask.sum(), + global_valid_toks=torch.tensor(2.0), + logits=student_logits, + student_logits_contig=student_logits, + teacher_full_logits_by_idx={0: teacher_logits}, + aligns_by_idx={0: align}, + ) + loss.backward() + return loss.detach(), metrics, student_logits.grad.detach() + + clean_loss, clean_metrics, clean_grad = run(teacher_clean) + corrupt_loss, corrupt_metrics, corrupt_grad = run(teacher_corrupt) + + torch.testing.assert_close(corrupt_loss, clean_loss, rtol=1e-6, atol=1e-6) + torch.testing.assert_close(corrupt_grad, clean_grad, rtol=1e-6, atol=1e-6) + assert clean_metrics["num_valid_samples"] == 1 + assert corrupt_metrics["num_valid_samples"] == 1 + assert torch.count_nonzero(clean_grad[1]).item() == 0 + assert torch.count_nonzero(clean_grad[0, 2:]).item() == 0 + + +def test_same_vocab_chat_kd_uses_semantic_mask_and_its_own_denominator(tmp_path): + """Template scaffold stays in CE but cannot enter direct same-vocab KD.""" + cfg = _ct_loss_cfg(str(tmp_path / "unused_projection.pt")) + cfg["projection_matrix_paths"] = [None] + cfg["teacher_vocab_sizes"] = [_CT_V_STUDENT] + cfg["vocab_topk"] = 3 + cfg["temperature"] = 1.0 + cfg["ce_loss_scale"] = 0.0 + loss_fn = CSDLossFn(cfg) + + torch.manual_seed(19) + student = torch.randn(1, 5, _CT_V_STUDENT, requires_grad=True) + teacher = torch.randn(1, 5, _CT_V_STUDENT) + input_ids = torch.tensor([[0, 1, 2, 3, 4]]) + ce_mask = torch.ones(1, 5, dtype=torch.long) + # Only labels at positions 3 (content) and 4 (explicit EOT) are KD targets, + # hence predictors 2 and 3 receive KD gradients. + kd_mask = torch.tensor([[0, 0, 0, 1, 1]]) + sample_mask = torch.ones(1) + data = BatchedDataDict( + { + "input_ids": input_ids, + "input_lengths": torch.tensor([5]), + "token_mask": ce_mask, + "kd_token_mask": kd_mask, + "sample_mask": sample_mask, + } + ) + align = LocalizedAlignment( + sample_mask=sample_mask, + student_input_ids=input_ids, + student_token_mask=ce_mask, + student_kd_token_mask=kd_mask, + ) + expected_kd = same_tokenizer_topk_kl( + student, + teacher, + kd_mask=kd_mask, + sample_mask=sample_mask, + num_kd_tokens=torch.tensor(2.0), + vocab_size=_CT_V_STUDENT, + vocab_topk=3, + temperature=1.0, + reverse_kl=loss_fn.reverse_kl, + tp_group=None, + cp_group=None, + ) + loss, metrics = loss_fn( + data=data, + global_valid_seqs=torch.tensor(1.0), + # CE has four next-token targets, while semantic KD has two. + global_valid_toks=torch.tensor(4.0), + global_valid_kd_toks=torch.tensor(2.0), + logits=student, + student_logits_contig=student, + teacher_full_logits_by_idx={0: teacher}, + aligns_by_idx={0: align}, + ) + torch.testing.assert_close(loss, expected_kd) + assert metrics["kl_loss"] == pytest.approx(expected_kd.item()) + loss.backward() + assert torch.count_nonzero(student.grad[0, :2]).item() == 0 + assert torch.count_nonzero(student.grad[0, 2:4]).item() > 0 + assert torch.count_nonzero(student.grad[0, 4]).item() == 0 + + +def _run_kd_ce_combination(tmp_path, cfg_overrides, *, kd_mask=None): + """Run the full loss on a fixed same-vocab batch; return loss, metrics, grad. + + The KD/CE combination does not depend on how the KD term was computed, so a + same-vocab teacher keeps the batch small. + """ + cfg = _ct_loss_cfg(str(tmp_path / "unused_projection.pt")) + cfg["projection_matrix_paths"] = [None] + cfg["teacher_vocab_sizes"] = [_CT_V_STUDENT] + cfg["vocab_topk"] = 3 + cfg["temperature"] = 1.0 + cfg.update(cfg_overrides) + loss_fn = CSDLossFn(cfg) + + torch.manual_seed(23) + student = torch.randn(1, 5, _CT_V_STUDENT, requires_grad=True) + teacher = torch.randn(1, 5, _CT_V_STUDENT) + input_ids = torch.tensor([[0, 1, 2, 3, 4]]) + token_mask = torch.ones(1, 5, dtype=torch.long) + if kd_mask is None: + kd_mask = token_mask + sample_mask = torch.ones(1) + data = BatchedDataDict( + { + "input_ids": input_ids, + "input_lengths": torch.tensor([5]), + "token_mask": token_mask, + "kd_token_mask": kd_mask, + "sample_mask": sample_mask, + } + ) + align = LocalizedAlignment( + sample_mask=sample_mask, + student_input_ids=input_ids, + student_token_mask=token_mask, + student_kd_token_mask=kd_mask, + ) + loss, metrics = loss_fn( + data=data, + global_valid_seqs=torch.tensor(1.0), + global_valid_toks=torch.tensor(4.0), + global_valid_kd_toks=torch.tensor(4.0), + logits=student, + student_logits_contig=student, + teacher_full_logits_by_idx={0: teacher}, + aligns_by_idx={0: align}, + ) + loss.backward() + return loss.detach(), metrics, student.grad.detach() + + +_KD_ONLY = {"dynamic_loss_scaling": False, "kl_loss_weight": 1.0, "ce_loss_scale": 0.0} +_CE_ONLY = {"dynamic_loss_scaling": False, "kl_loss_weight": 0.0, "ce_loss_scale": 1.0} + + +def test_csd_fixed_weighting_scales_kd_and_ce(tmp_path): + """Fixed weighting is kl_loss_weight * KD + ce_loss_scale * CE in value and gradient.""" + _, _, kd_grad = _run_kd_ce_combination(tmp_path, _KD_ONLY) + _, _, ce_grad = _run_kd_ce_combination(tmp_path, _CE_ONLY) + loss, metrics, grad = _run_kd_ce_combination( + tmp_path, + {"dynamic_loss_scaling": False, "kl_loss_weight": 5.0, "ce_loss_scale": 0.3}, + ) + + kd, ce = metrics["kl_loss"], metrics["ce_loss"] + assert kd > 0 and ce > 0 + assert loss.item() == pytest.approx(5.0 * kd + 0.3 * ce) + assert metrics["kl_loss_scale"] == 1.0 + torch.testing.assert_close(grad, 5.0 * kd_grad + 0.3 * ce_grad) + + +def test_csd_dynamic_scaling_matches_kd_to_ce_and_ignores_weights(tmp_path): + """Dynamic scaling is sg(CE / KD) * KD + CE; the fixed weights have no effect.""" + _, _, kd_grad = _run_kd_ce_combination(tmp_path, _KD_ONLY) + _, _, ce_grad = _run_kd_ce_combination(tmp_path, _CE_ONLY) + loss, metrics, grad = _run_kd_ce_combination( + tmp_path, + {"dynamic_loss_scaling": True, "kl_loss_weight": 1.0, "ce_loss_scale": 0.1}, + ) + reweighted_loss, _, reweighted_grad = _run_kd_ce_combination( + tmp_path, + {"dynamic_loss_scaling": True, "kl_loss_weight": 5.0, "ce_loss_scale": 1.0}, + ) + + kd, ce = metrics["kl_loss"], metrics["ce_loss"] + assert kd > 0 and ce > 0 + assert metrics["kl_loss_scale"] == pytest.approx(ce / kd) + # The rescaled KD term equals CE in value, so the loss reads as 2 * CE. + assert loss.item() == pytest.approx(2 * ce) + torch.testing.assert_close(grad, (ce / kd) * kd_grad + ce_grad) + torch.testing.assert_close(reweighted_loss, loss) + torch.testing.assert_close(reweighted_grad, grad) + + +def test_csd_dynamic_scaling_keeps_unit_scale_when_kd_is_zero(tmp_path): + """A microbatch without KD positions keeps scale 1, so the loss is CE alone.""" + loss, metrics, grad = _run_kd_ce_combination( + tmp_path, + {"dynamic_loss_scaling": True}, + kd_mask=torch.zeros(1, 5, dtype=torch.long), + ) + _, _, ce_grad = _run_kd_ce_combination(tmp_path, _CE_ONLY) + + assert metrics["kl_loss"] == 0.0 + assert metrics["kl_loss_scale"] == 1.0 + assert loss.item() == pytest.approx(metrics["ce_loss"]) + torch.testing.assert_close(grad, ce_grad) + + +def test_csd_table_paths_identify_cross_tokenizer_without_projection(tmp_path): + """CSD table artifacts select the cross-tokenizer path without dummy matrices.""" + cfg = _ct_loss_cfg(None) + cfg["student_to_teacher_table_paths"] = [str(tmp_path / "forward.pt")] + cfg["teacher_to_student_table_paths"] = [str(tmp_path / "reverse.pt")] + cfg["common_indices_from_subtoks"] = True + loss_fn = CSDLossFn(cfg) + assert not loss_fn._teacher_is_same_vocab(0) + assert loss_fn.projection_matrix_paths == [None] + + +def test_csd_runtime_table_lists_must_match_teachers(tmp_path): + cfg = _ct_loss_cfg(None) + cfg["student_to_teacher_table_paths"] = [None, None] + with pytest.raises(ValueError, match="one entry per teacher"): + CSDLossFn(cfg) + + +def _same_vocab_megatron_case(): + torch.manual_seed(5) + cfg = _ct_loss_cfg(None) + cfg.update( + projection_matrix_paths=[None], + teacher_vocab_sizes=[_CT_V_STUDENT], + vocab_topk=4, + dynamic_loss_scaling=True, + ) + input_ids = torch.randint(0, _CT_V_STUDENT, (2, 8)) + token_mask = torch.ones((2, 8), dtype=torch.long) + token_mask[1, 6:] = 0 + data = BatchedDataDict( + { + "input_ids": input_ids, + "input_lengths": torch.tensor([8, 6]), + "token_mask": token_mask, + "kd_token_mask": token_mask, + "sample_mask": torch.tensor([1, 1]), + } + ) + valid_toks = (token_mask[:, 1:]).sum().float() + return ( + cfg, + data, + valid_toks, + torch.randn(2, 8, _CT_V_STUDENT), + torch.randn(2, 8, _CT_V_STUDENT), + ) + + +def _replicated_same_vocab_loss( + loss_fn, data, valid_toks, student, teacher, *, cp_group, tp_group, window +): + contiguous = cp_load_balanced_to_contiguous(student, cp_group=cp_group) + align = LocalizedAlignment( + sample_mask=data["sample_mask"], + student_input_ids=data["input_ids"][:, window], + student_token_mask=data["token_mask"][:, window], + student_kd_token_mask=data["kd_token_mask"][:, window], + ) + return loss_fn( + data, + data["sample_mask"].sum(), + valid_toks, + student, + contiguous, + {0: teacher[:, window]}, + {0: align}, + global_valid_kd_toks=valid_toks, + tp_group=tp_group, + cp_group=cp_group, + cp_normalized_by_caller=True, + ) + + +def _run_same_vocab_megatron_cp(rank: int, world_size: int, init_file: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + backend="gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + ) + try: + cp_group = dist.new_group(ranks=list(range(world_size))) + tp_groups = [dist.new_group(ranks=[r]) for r in range(world_size)] + cfg, data, valid_toks, student_full, teacher = _same_vocab_megatron_case() + loss_fn = CSDLossFn(cfg) + + reference = student_full.clone().requires_grad_(True) + ref_loss, ref_metrics = loss_fn( + data, + data["sample_mask"].sum(), + valid_toks, + reference, + reference, + {0: teacher}, + { + 0: LocalizedAlignment( + sample_mask=data["sample_mask"], + student_input_ids=data["input_ids"], + student_token_mask=data["token_mask"], + student_kd_token_mask=data["kd_token_mask"], + ) + }, + global_valid_kd_toks=valid_toks, + ) + ref_loss.backward() + + positions = _get_tokens_on_this_cp_rank( + torch.arange(8), rank, world_size, seq_dim=0 + ) + local = student_full[:, positions].clone().requires_grad_(True) + width = 8 // world_size + loss, metrics = _replicated_same_vocab_loss( + loss_fn, + data, + valid_toks, + local, + teacher, + cp_group=cp_group, + tp_group=tp_groups[rank], + window=slice(rank * width, (rank + 1) * width), + ) + # Megatron's replicated path divides the loss by the CP size. + (loss / world_size).backward() + + torch.testing.assert_close( + loss.detach(), ref_loss.detach(), rtol=1e-5, atol=1e-6 + ) + for key in ("kl_loss", "ce_loss", "kl_loss_scale"): + assert metrics[key] == pytest.approx(ref_metrics[key], rel=1e-5), key + gathered = [torch.empty_like(local.grad) for _ in range(world_size)] + dist.all_gather(gathered, local.grad) + full_grad = torch.zeros_like(student_full) + for cp_rank, shard in enumerate(gathered): + full_grad[ + :, + _get_tokens_on_this_cp_rank( + torch.arange(8), cp_rank, world_size, seq_dim=0 + ), + ] = shard + torch.testing.assert_close(full_grad, reference.grad, rtol=1e-5, atol=1e-6) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_gloo_available(), reason="gloo backend unavailable") +def test_same_vocab_kd_on_megatron_cp2_matches_cp1(tmp_path): + """The per-window same-tokenizer KD keeps full weight under Megatron's /CP.""" + mp.spawn( + _run_same_vocab_megatron_cp, + args=(2, os.fspath(tmp_path / "gloo_init")), + nprocs=2, + join=True, + ) diff --git a/tests/unit/algorithms/csd/test_csd_loss_parity.py b/tests/unit/algorithms/csd/test_csd_loss_parity.py new file mode 100644 index 00000000000..8c50d2a4792 --- /dev/null +++ b/tests/unit/algorithms/csd/test_csd_loss_parity.py @@ -0,0 +1,895 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU regression coverage for CSD values, gradients, masking, and transport.""" + +import os +import tempfile + +import pytest +import torch + + +def _write_subtok_tables(tmpdir, v_s, v_t): + """Tiny forward/reverse pseudo-target tables. + + Student/teacher ids 0..9 are 1-to-1 common (length-1 chain, mapped to the + same id). To exercise the mismatch (prefix-support) path with + ``keep_realized=False`` (which needs >= 2 support pairs per chunk): + + * Forward (student->teacher): students 10 and 11 both have a length-2 + chain with teacher prefix ``(8,)`` and distinct finals -> the shared + key ``(2, (8,))`` maps to two ``(s, t)`` pairs, so a 1-to-2 chunk whose + teacher prefix is ``8`` finds 2 support pairs. + * Reverse (teacher->student): teachers 10 and 11 both have a length-2 + chain with student prefix ``(9,)`` and distinct finals -> the shared + key ``(2, (9,))`` maps to two pairs, so a 2-to-1 chunk whose student + prefix is ``9`` finds 2 support pairs. + """ + max_chain = 3 + subtoks = torch.full((v_s, max_chain), -1, dtype=torch.long) + lengths = torch.zeros((v_s,), dtype=torch.long) + for s in range(10): + subtoks[s, 0] = s + lengths[s] = 1 + # Shared forward prefix (8,) with two distinct finals. + subtoks[10, 0], subtoks[10, 1], lengths[10] = 8, 3, 2 + subtoks[11, 0], subtoks[11, 1], lengths[11] = 8, 4, 2 + # Filler length-2 chains (distinct prefixes; harmless). + for s in range(12, v_s): + subtoks[s, 0], subtoks[s, 1], lengths[s] = s % 10, (s + 3) % 10, 2 + fwd_path = os.path.join(tmpdir, "fwd_subtoks.pt") + torch.save({"subtoks": subtoks, "lengths": lengths}, fwd_path) + + subtoks_r = torch.full((v_t, max_chain), -1, dtype=torch.long) + lengths_r = torch.zeros((v_t,), dtype=torch.long) + for t in range(10): + subtoks_r[t, 0] = t + lengths_r[t] = 1 + # Shared reverse prefix (9,) with two distinct finals. + subtoks_r[10, 0], subtoks_r[10, 1], lengths_r[10] = 9, 3, 2 + subtoks_r[11, 0], subtoks_r[11, 1], lengths_r[11] = 9, 4, 2 + for t in range(12, v_t): + subtoks_r[t, 0], subtoks_r[t, 1], lengths_r[t] = t % 10, (t + 3) % 10, 2 + rev_path = os.path.join(tmpdir, "rev_subtoks.pt") + torch.save({"subtoks": subtoks_r, "lengths": lengths_r}, rev_path) + return fwd_path, rev_path + + +def _csd_knobs(fwd_path, rev_path, v_t): + """CSD settings shared by dense and sparse transport checks.""" + return { + "temperature": 1.0, + "teacher_vocab_size": v_t, + "common_indices_from_subtoks": True, + "student_to_teacher_table_path": fwd_path, + "teacher_to_student_table_path": rev_path, + "kl_chunk_shift": True, + "csd_loss_fn": "jsd", + "csd_last_pos_loss_fn": "kl", + "csd_jsd_beta": 0.5, + "csd_mismatch_loss_beta": 2.0, + "csd_noise_filter_topk": 0, + "reverse_kl": False, + } + + +def _build_case(name, tmpdir, seed=0): + """Return a serializable fixture dict for the named case. + + Cases: + * ``common_only``: three 1-to-1 common chunks per sample (exercises the + common-vocab JSD partition KL, no mismatch/prefix-index path). + * ``with_mismatch``: adds a 1-to-2 and a 2-to-1 chunk (exercises the + prefix support index + batched mismatch partition KL). + * ``mismatch_only``: a 1-to-2 chunk at the sequence start (never + shifted) followed by a shifted 2-to-1 chunk, with no common chunk. + """ + torch.manual_seed(seed) + v_s, v_t = 16, 16 + fwd_path, rev_path = _write_subtok_tables(tmpdir, v_s, v_t) + B = 2 + + if name == "common_only": + S = T = 4 + # Every position carries a common id (0..9). + student_ids = torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.long) + teacher_ids = student_ids.clone() + chunks = [(0, 1, 0, 1), (1, 2, 1, 2), (2, 3, 2, 3)] # (ss, se, ts, te) + per_sample_chunks = [chunks, chunks] + elif name == "with_mismatch": + S = T = 6 + # chunk 0: 1-to-1 common (id in 0..9). chunk 1: 1-to-2 (M=1,N=2) whose + # teacher prefix is 8 -> forward key (2,(8,)) has 2 support pairs. + # chunk 2: 2-to-1 (M=2,N=1) whose student prefix is 9 -> reverse key + # (2,(9,)) has 2 support pairs. Positions 2 and 5 are gaps. + student_ids = torch.tensor( + [[0, 5, 0, 9, 5, 0], [1, 6, 0, 9, 6, 0]], dtype=torch.long + ) + teacher_ids = torch.tensor( + [[0, 8, 3, 0, 5, 0], [1, 8, 4, 0, 6, 0]], dtype=torch.long + ) + chunks = [(0, 1, 0, 1), (1, 2, 1, 3), (3, 5, 4, 5)] + per_sample_chunks = [chunks, chunks] + elif name == "mismatch_only": + S = T = 3 + student_ids = torch.tensor([[5, 9, 5], [6, 9, 6]], dtype=torch.long) + teacher_ids = torch.tensor([[8, 3, 5], [8, 4, 6]], dtype=torch.long) + chunks = [(0, 1, 0, 2), (1, 3, 2, 3)] + per_sample_chunks = [chunks, chunks] + else: + raise ValueError(name) + + max_pairs = max(len(c) for c in per_sample_chunks) + s_spans = torch.zeros(B, max_pairs, 2, dtype=torch.long) + t_spans = torch.zeros(B, max_pairs, 2, dtype=torch.long) + pair_valid = torch.zeros(B, max_pairs, dtype=torch.bool) + num_chunks = torch.zeros(B, dtype=torch.long) + for b in range(B): + cs = per_sample_chunks[b] + num_chunks[b] = len(cs) + for k, (ss, se, ts, te) in enumerate(cs): + s_spans[b, k, 0], s_spans[b, k, 1] = ss, se + t_spans[b, k, 0], t_spans[b, k, 1] = ts, te + pair_valid[b, k] = True + + student_logits = torch.randn(B, S, v_s, dtype=torch.float32) + teacher_logits = torch.randn(B, T, v_t, dtype=torch.float32) + # global_valid_chunks = total valid chunks across the (single-rank) batch. + gvc = float(sum(len(c) for c in per_sample_chunks)) + + return { + "v_s": v_s, + "v_t": v_t, + "student_logits": student_logits, + "teacher_logits": teacher_logits, + "student_ids": student_ids, + "teacher_ids": teacher_ids, + "s_spans": s_spans, + "t_spans": t_spans, + "pair_valid": pair_valid, + "num_chunks": num_chunks, + "global_valid_chunks": gvc, + "knobs": _csd_knobs(fwd_path, rev_path, v_t), + } + + +# ----------------------------------------------------------------------------- # +# In-process CSD loss runner. +# ----------------------------------------------------------------------------- # +def _run_csd( + fx, + *, + teacher_sparse_payload=None, + sample_mask=None, + global_valid_chunks=None, + cfg_overrides=None, +): + from nemo_rl.algorithms.loss.csd import ( + CSDLossFn, + ) + from nemo_rl.algorithms.x_token.loss_utils import LocalizedAlignment + + v_s, v_t = fx["v_s"], fx["v_t"] + knobs = fx["knobs"] + cfg = { + "temperature": knobs["temperature"], + "vocab_topk": 8, # same-vocab path only; unused here (cross-tok teacher) + "reverse_kl": knobs["reverse_kl"], + "kl_loss_weight": 1.0, + "ce_loss_scale": 1.0, + "dynamic_loss_scaling": False, + "kd_loss_mode": "sum", + "normalize_teacher_by_vocab": False, + "alpha": 1.0, + "sum_weights_metric": None, + "student_vocab_size": v_s, + "teacher_vocab_sizes": [v_t], + "projection_matrix_paths": [None], + "teacher_weights": [1.0], + "common_indices_from_subtoks": knobs["common_indices_from_subtoks"], + "student_to_teacher_table_paths": [knobs["student_to_teacher_table_path"]], + "teacher_to_student_table_paths": [knobs["teacher_to_student_table_path"]], + "kl_chunk_shift": knobs["kl_chunk_shift"], + "csd_loss_fn": knobs["csd_loss_fn"], + "csd_last_pos_loss_fn": knobs["csd_last_pos_loss_fn"], + "csd_jsd_beta": knobs["csd_jsd_beta"], + "csd_mismatch_loss_beta": knobs["csd_mismatch_loss_beta"], + "csd_noise_filter_topk": knobs["csd_noise_filter_topk"], + "teacher_topk_ipc_k": v_t if teacher_sparse_payload is not None else 0, + "teacher_topk_ipc_support_mode": "row_topk", + # Full support needs no forced insertion. Keeping this false also makes + # dense and sparse mismatch eligibility identical for this regression. + "teacher_topk_ipc_keep_realized": False, + } + cfg.update(cfg_overrides or {}) + loss_fn = CSDLossFn(cfg) + student_logits = fx["student_logits"].clone().requires_grad_(True) + if sample_mask is None: + sample_mask = torch.ones(fx["student_ids"].shape[0], dtype=torch.bool) + if global_valid_chunks is None: + global_valid_chunks = fx["global_valid_chunks"] + align = LocalizedAlignment( + sample_mask=sample_mask, + pair_valid=fx["pair_valid"], + student_input_ids=fx["student_ids"], + teacher_input_ids=fx["teacher_ids"], + student_spans=fx["s_spans"], + teacher_spans=fx["t_spans"], + num_chunks=fx["num_chunks"], + ) + loss, metrics = loss_fn._compute_csd_loss( + 0, + student_logits, + None if teacher_sparse_payload is not None else fx["teacher_logits"].clone(), + align, + teacher_sparse_payload=teacher_sparse_payload, + teacher_vocab_size=v_t, + global_valid_chunks=torch.as_tensor(global_valid_chunks, dtype=torch.float32), + ) + loss.backward() + return loss.detach(), metrics, student_logits.grad.detach() + + +def test_csd_sparse_full_support_matches_dense(): + """The production sparse consumer matches dense when support is complete.""" + with tempfile.TemporaryDirectory() as tmp: + fx = _build_case("with_mismatch", tmp) + dense_loss, dense_metrics, dense_grad = _run_csd(fx) + + teacher_logits = fx["teacher_logits"].clone() + token_ids = torch.arange(fx["v_t"], dtype=torch.int32).view(1, 1, -1) + token_ids = token_ids.expand(*teacher_logits.shape).clone() + sparse_payload = ( + teacher_logits, + token_ids, + torch.logsumexp(teacher_logits / fx["knobs"]["temperature"], dim=-1), + None, + ) + sparse_loss, sparse_metrics, sparse_grad = _run_csd( + fx, + teacher_sparse_payload=sparse_payload, + ) + + torch.testing.assert_close(sparse_loss, dense_loss, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(sparse_grad, dense_grad, rtol=1e-5, atol=1e-6) + assert sparse_metrics["num_common_chunks"] == dense_metrics["num_common_chunks"] + assert sparse_metrics["num_mismatch_chunks"] == dense_metrics["num_mismatch_chunks"] + + +def test_csd_dense_teacher_stays_cp_local_and_slices_padded_vocab(monkeypatch): + """Dense CSD never sequence-gathers teacher logits and excludes padded cols.""" + from nemo_rl.algorithms.loss import csd as loss_functions_module + + with tempfile.TemporaryDirectory() as tmp: + fx = _build_case("with_mismatch", tmp) + baseline_loss, baseline_metrics, baseline_grad = _run_csd(fx) + + padded_fx = dict(fx) + padded_fx["teacher_logits"] = torch.cat( + [ + fx["teacher_logits"], + torch.full( + (*fx["teacher_logits"].shape[:-1], 3), + 1.0e4, + dtype=fx["teacher_logits"].dtype, + ), + ], + dim=-1, + ) + original_gather = loss_functions_module.allgather_cp_contiguous_tensor + + def reject_dense_teacher_gather(tensor, group, seq_dim=1): + if ( + tensor.ndim == 3 + and not tensor.requires_grad + and tensor.shape[-1] >= fx["v_t"] + ): + raise AssertionError( + "dense teacher logits entered a full-sequence CP gather" + ) + return original_gather(tensor, group, seq_dim) + + monkeypatch.setattr( + loss_functions_module, + "allgather_cp_contiguous_tensor", + reject_dense_teacher_gather, + ) + padded_loss, padded_metrics, padded_grad = _run_csd(padded_fx) + + torch.testing.assert_close(padded_loss, baseline_loss, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(padded_grad, baseline_grad, rtol=1e-5, atol=1e-6) + for key in ( + "kl_common_per_chunk", + "kl_partition_last_per_chunk", + "kl_mismatch_scaled_per_chunk", + "top1_acc_per_chunk", + "num_common_chunks", + "num_mismatch_chunks", + ): + assert padded_metrics[key] == pytest.approx(baseline_metrics[key]) + + +def test_dense_teacher_lookup_broadcasts_repeated_out_of_order_queries(): + """The CP1 lookup seam preserves arbitrary advanced-index semantics.""" + from nemo_rl.algorithms.loss.csd import ( + CSDLossFn, + ) + + teacher = torch.tensor( + [ + [[-4.0, 0.0, 2.0], [3.0, -2.0, 1.0]], + [[7.0, 5.0, -1.0], [0.0, -6.0, 8.0]], + ] + ) + batch = torch.tensor([[1], [0], [1]]) + positions = torch.tensor([[1], [0], [1]]) + token_ids = torch.tensor([[2, 0, 2], [1, 2, 1], [0, 1, 0]]) + actual = CSDLossFn._lookup_cp_sharded_teacher_logits( + teacher, + batch, + positions, + token_ids, + cp_group=None, + ) + b_full, p_full, t_full = torch.broadcast_tensors(batch, positions, token_ids) + expected = teacher[b_full, p_full, t_full] + torch.testing.assert_close(actual, expected) + + +def test_student_lookup_tp1_cp1_preserves_repeated_query_gradients(): + """The owner-lookup fast path matches advanced indexing and its gradient.""" + from nemo_rl.algorithms.loss.csd import ( + CSDLossFn, + ) + + student = torch.randn(2, 3, 5, requires_grad=True) + reference = student.detach().clone().requires_grad_(True) + batch = torch.tensor([[1], [0], [1]]) + positions = torch.tensor([[2], [0], [2]]) + token_ids = torch.tensor([[4, 0, 4], [1, 3, 1], [2, 2, 2]]) + weights = torch.arange(1, 10, dtype=student.dtype).view(3, 3) + + actual = CSDLossFn._lookup_tp_cp_sharded_student_logits( + student, + batch, + positions, + token_ids, + tp_group=None, + cp_group=None, + ) + b_full, p_full, t_full = torch.broadcast_tensors(batch, positions, token_ids) + expected = reference[b_full, p_full, t_full] + torch.testing.assert_close(actual, expected) + + (actual * weights).sum().backward() + (expected * weights).sum().backward() + torch.testing.assert_close(student.grad, reference.grad) + + +def test_student_lookup_tp2_cp2_preserves_outer_product_gradients(monkeypatch): + """Compact broadcast queries preserve owner masking and local gradients.""" + from nemo_rl.algorithms.loss import csd as loss_functions_module + + tp_group = object() + cp_group = object() + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group: 2) + monkeypatch.setattr( + torch.distributed, + "get_rank", + lambda group: 1 if group is tp_group else 0, + ) + collective_calls = [] + + def tp_reduce(values, group): + collective_calls.append(("tp", group)) + return values + + def cp_reduce(values, group): + collective_calls.append(("cp", group)) + return values + + monkeypatch.setattr( + loss_functions_module, + "group_all_reduce_sum_with_grad", + tp_reduce, + ) + monkeypatch.setattr( + loss_functions_module, + "group_all_reduce_sum_with_grad_backward_sum", + cp_reduce, + ) + + student = torch.randn(2, 3, 5, requires_grad=True) + reference = student.detach().clone().requires_grad_(True) + batch = torch.tensor([[1], [0], [1]]) + positions = torch.tensor([[0], [4], [2]]) + token_ids = torch.tensor([[0, 6, 9, 4]]) + weights = torch.arange(1, 13, dtype=student.dtype).view(3, 4) + + actual = loss_functions_module.CSDLossFn._lookup_tp_cp_sharded_student_logits( + student, + batch, + positions, + token_ids, + tp_group=tp_group, + cp_group=cp_group, + ) + expected = reference[ + batch, + positions.remainder(reference.shape[1]), + token_ids.remainder(reference.shape[2]), + ] + expected = torch.where(positions // reference.shape[1] == 0, expected, 0.0) + expected = torch.where(token_ids // reference.shape[2] == 1, expected, 0.0) + + torch.testing.assert_close(actual, expected) + (actual * weights).sum().backward() + (expected * weights).sum().backward() + torch.testing.assert_close(student.grad, reference.grad) + assert collective_calls == [("tp", tp_group), ("cp", cp_group)] + + +def test_student_lookup_dense_common_query_saves_compact_indices(monkeypatch): + """The qualification-sized owner lookup retains no dense index storage.""" + from nemo_rl.algorithms.loss import csd as loss_functions_module + + tp_group = object() + cp_group = object() + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group: 2) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group: 0) + collective_calls = [] + + def tp_reduce(values, group): + collective_calls.append(("tp", group)) + return values + + def cp_reduce(values, group): + collective_calls.append(("cp", group)) + return values + + monkeypatch.setattr( + loss_functions_module, + "group_all_reduce_sum_with_grad", + tp_reduce, + ) + monkeypatch.setattr( + loss_functions_module, + "group_all_reduce_sum_with_grad_backward_sum", + cp_reduce, + ) + + row_count = 64 + common_width = 109_567 + student = torch.arange(8, dtype=torch.float32).view(1, 2, 4) + student.requires_grad_(True) + batch = torch.zeros((row_count, 1), dtype=torch.long) + # Rank CP0/TP0 owns none of these queries. It must still execute both + # collectives, and the autograd graph must retain only the compact axes. + positions = torch.full((row_count, 1), 2, dtype=torch.long) + token_ids = (torch.arange(common_width) % 4 + 4).view(1, common_width) + saved_storage = [] + + def save_tensor(tensor): + if tensor.dtype in (torch.long, torch.bool): + saved_storage.append((tensor.dtype, tensor.untyped_storage().nbytes())) + return tensor + + with torch.autograd.graph.saved_tensors_hooks(save_tensor, lambda tensor: tensor): + actual = loss_functions_module.CSDLossFn._lookup_tp_cp_sharded_student_logits( + student, + batch, + positions, + token_ids, + tp_group=tp_group, + cp_group=cp_group, + ) + actual.sum().backward() + + assert actual.shape == (row_count, common_width) + assert torch.count_nonzero(actual) == 0 + assert torch.count_nonzero(student.grad) == 0 + assert collective_calls == [("tp", tp_group), ("cp", cp_group)] + long_storage = sorted(size for dtype, size in saved_storage if dtype == torch.long) + bool_storage = sorted(size for dtype, size in saved_storage if dtype == torch.bool) + assert long_storage == [row_count * 8, row_count * 8, common_width * 8] + assert bool_storage == [row_count, common_width] + + +@pytest.mark.parametrize( + ("batch", "position", "token", "match"), + [ + (torch.tensor([-1]), torch.tensor([0]), torch.tensor([0]), "batch"), + (torch.tensor([0]), torch.tensor([3]), torch.tensor([0]), "sequence"), + (torch.tensor([0]), torch.tensor([0]), torch.tensor([5]), "vocabulary"), + ], +) +def test_student_lookup_rejects_queries_before_collectives( + batch: torch.Tensor, + position: torch.Tensor, + token: torch.Tensor, + match: str, +): + """Malformed globally replicated owner queries fail before communication.""" + from nemo_rl.algorithms.loss.csd import ( + CSDLossFn, + ) + + with pytest.raises(IndexError, match=match): + CSDLossFn._lookup_tp_cp_sharded_student_logits( + torch.zeros(1, 3, 5), + batch, + position, + token, + tp_group=None, + cp_group=None, + ) + + +def test_student_lookup_rejects_nonbroadcastable_queries_before_collectives( + monkeypatch, +): + """Incompatible compact query axes fail before either collective.""" + from nemo_rl.algorithms.loss import csd as loss_functions_module + + collective_calls = [] + + def unexpected_collective(values, group): + collective_calls.append(group) + return values + + monkeypatch.setattr( + loss_functions_module, + "group_all_reduce_sum_with_grad", + unexpected_collective, + ) + monkeypatch.setattr( + loss_functions_module, + "group_all_reduce_sum_with_grad_backward_sum", + unexpected_collective, + ) + + with pytest.raises(RuntimeError): + loss_functions_module.CSDLossFn._lookup_tp_cp_sharded_student_logits( + torch.zeros(1, 3, 5), + torch.zeros(2, dtype=torch.long), + torch.zeros(3, dtype=torch.long), + torch.zeros(2, dtype=torch.long), + tp_group=object(), + cp_group=object(), + ) + + assert collective_calls == [] + + +def test_csd_student_logits_never_enter_rank3_gathers(monkeypatch): + """The CSD student path gathers only scalar log-normalizers across CP.""" + from nemo_rl.algorithms.loss import csd as loss_functions_module + + with tempfile.TemporaryDirectory() as tmp: + fx = _build_case("with_mismatch", tmp) + baseline_loss, baseline_metrics, baseline_grad = _run_csd(fx) + original_gather = loss_functions_module.allgather_cp_contiguous_tensor + + def reject_rank3_gather(tensor, group, seq_dim=1): + if tensor.ndim == 3: + raise AssertionError("CSD attempted a rank-3 CP gather") + return original_gather(tensor, group, seq_dim) + + monkeypatch.setattr( + loss_functions_module, + "allgather_cp_contiguous_tensor", + reject_rank3_gather, + ) + loss, metrics, grad = _run_csd(fx) + + torch.testing.assert_close(loss, baseline_loss, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(grad, baseline_grad, rtol=1e-5, atol=1e-6) + for key in ( + "kl_common_per_chunk", + "kl_partition_last_per_chunk", + "top1_acc_per_chunk", + "num_common_chunks", + "num_mismatch_chunks", + ): + assert metrics[key] == pytest.approx(baseline_metrics[key]) + + +def _single_sample_case(fx, sample_idx): + """Slice the batched fixture while retaining its on-disk CSD tables/config.""" + single = dict(fx) + for key in ( + "student_logits", + "teacher_logits", + "student_ids", + "teacher_ids", + "s_spans", + "t_spans", + "pair_valid", + "num_chunks", + ): + single[key] = fx[key][sample_idx : sample_idx + 1].clone() + single["global_valid_chunks"] = float( + single["pair_valid"][0, : single["num_chunks"][0]].sum().item() + ) + return single + + +@pytest.mark.parametrize("case_name", ["common_only", "with_mismatch"]) +def test_csd_sample_mask_matches_single_sample_and_zeroes_masked_grad(case_name): + """Masked rows affect neither CSD cohorts nor forward/backward numerators.""" + with tempfile.TemporaryDirectory() as tmp: + fx = _build_case(case_name, tmp) + single = _single_sample_case(fx, 0) + single_loss, single_metrics, single_grad = _run_csd(single) + + sample_mask = torch.tensor([True, False]) + masked_loss, masked_metrics, masked_grad = _run_csd( + fx, + sample_mask=sample_mask, + global_valid_chunks=single["global_valid_chunks"], + ) + + # Corrupt every logit in the masked row. It must not change selected + # cohorts, scalar numerators, or the valid row's gradient. + corrupted = dict(fx) + corrupted["student_logits"] = fx["student_logits"].clone() + corrupted["teacher_logits"] = fx["teacher_logits"].clone() + corrupted["student_logits"][1] = torch.linspace( + -50.0, + 50.0, + corrupted["student_logits"][1].numel(), + ).reshape_as(corrupted["student_logits"][1]) + corrupted["teacher_logits"][1] = torch.linspace( + 75.0, + -75.0, + corrupted["teacher_logits"][1].numel(), + ).reshape_as(corrupted["teacher_logits"][1]) + corrupt_loss, corrupt_metrics, corrupt_grad = _run_csd( + corrupted, + sample_mask=sample_mask, + global_valid_chunks=single["global_valid_chunks"], + ) + + torch.testing.assert_close(masked_loss, single_loss, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(masked_grad[0], single_grad[0], rtol=1e-5, atol=1e-6) + assert torch.count_nonzero(masked_grad[1]).item() == 0 + torch.testing.assert_close(corrupt_loss, masked_loss, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(corrupt_grad, masked_grad, rtol=1e-5, atol=1e-6) + + for key in ( + "kl_common_per_chunk", + "kl_partition_last_per_chunk", + "top1_acc_per_chunk", + "num_common_chunks", + "num_mismatch_chunks", + "num_noise_filtered_common_chunks", + "num_noise_filtered_mismatch_chunks", + "num_valid_samples", + ): + assert masked_metrics[key] == pytest.approx(single_metrics[key]) + assert corrupt_metrics[key] == pytest.approx(masked_metrics[key]) + + +@pytest.mark.parametrize("case_name", ["common_only", "with_mismatch"]) +def test_csd_all_samples_masked_is_zero_safe(case_name): + """An empty CSD cohort returns connected zero with an exactly zero gradient.""" + with tempfile.TemporaryDirectory() as tmp: + fx = _build_case(case_name, tmp) + loss, metrics, grad = _run_csd( + fx, + sample_mask=torch.zeros(2, dtype=torch.bool), + global_valid_chunks=0.0, + ) + + assert torch.isfinite(loss) + assert torch.equal(loss, torch.zeros_like(loss)) + assert torch.count_nonzero(grad).item() == 0 + for key in ( + "kl_common_per_chunk", + "kl_partition_last_per_chunk", + "top1_acc_per_chunk", + "num_common_chunks", + "num_mismatch_chunks", + "num_noise_filtered_common_chunks", + "num_noise_filtered_mismatch_chunks", + "num_valid_samples", + ): + assert metrics[key] == pytest.approx(0.0) + + +def test_csd_common_partition_matches_probability_reference(tmp_path): + """The common-token objective and gradient match an independent JSD equation.""" + fx = _build_case("common_only", str(tmp_path)) + loss, metrics, grad = _run_csd(fx) + student = fx["student_logits"].clone().requires_grad_(True) + # Shifted predictor positions for the three fixture chunks are 0, 0, 1. + positions = torch.tensor([0, 0, 1]) + student_prob = student.softmax(-1)[:, positions, :] + teacher_prob = fx["teacher_logits"].softmax(-1)[:, positions, :] + student_partition = torch.cat( + [student_prob[..., :10], student_prob[..., 10:].sum(-1, keepdim=True)], + dim=-1, + ) + teacher_partition = torch.cat( + [teacher_prob[..., :10], teacher_prob[..., 10:].sum(-1, keepdim=True)], + dim=-1, + ) + mixture = (student_partition + teacher_partition) / 2 + reference = ( + ( + 0.5 * student_partition * (student_partition.log() - mixture.log()) + + 0.5 * teacher_partition * (teacher_partition.log() - mixture.log()) + ) + .sum(-1) + .mean() + ) + reference.backward() + torch.testing.assert_close(loss, reference, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(grad, student.grad, rtol=1e-5, atol=1e-6) + assert metrics["num_common_chunks"] == 6 + assert metrics["num_mismatch_chunks"] == 0 + + +def test_csd_mismatch_partition_aggregation(tmp_path): + """Forward/reverse mismatch cohorts keep their beta weight and global counts.""" + fx = _build_case("with_mismatch", str(tmp_path)) + loss, metrics, grad = _run_csd(fx) + assert metrics["num_common_chunks"] == 2 + assert metrics["num_mismatch_chunks"] == 4 + weighted_mismatch = 2 * metrics["kl_partition_last_per_chunk"] + assert metrics["kl_mismatch_scaled_per_chunk"] == pytest.approx(weighted_mismatch) + expected = (2 * metrics["kl_common_per_chunk"] + 4 * weighted_mismatch) / fx[ + "global_valid_chunks" + ] + assert loss.item() == pytest.approx(expected) + assert torch.isfinite(grad).all() + assert torch.count_nonzero(grad) > 0 + + +_REALIZED_ONLY = {"csd_realized_only_mismatch": True} + + +def _realized_chain_reference(fx, student_logits): + """Independent realized-chain-vs-REST KL for every mismatch chunk.""" + knobs = fx["knobs"] + temperature = knobs["temperature"] + student_logp = torch.log_softmax(student_logits / temperature, dim=-1) + teacher_logp = torch.log_softmax(fx["teacher_logits"] / temperature, dim=-1) + + def positions(start, end, shift): + if shift: + return list(range(start - 1, end - 1)), list(range(start, end)) + return list(range(start, end)), list(range(start, end)) + + def chain(logp, ids, b, preds, labels): + return sum(logp[b, p, ids[b, lab]] for p, lab in zip(preds, labels)) + + def with_rest(chain_logp): + rest = (1.0 - chain_logp.exp()).clamp_min(1e-12).log() + return torch.stack([chain_logp, rest]) + + total = torch.zeros(()) + for b in range(fx["student_ids"].shape[0]): + for k in range(int(fx["num_chunks"][b])): + ss, se = fx["s_spans"][b, k].tolist() + ts, te = fx["t_spans"][b, k].tolist() + shift = knobs["kl_chunk_shift"] and ss > 0 and ts > 0 + s_pred, s_lab = positions(ss, se, shift) + t_pred, t_lab = positions(ts, te, shift) + student = with_rest( + chain(student_logp, fx["student_ids"], b, s_pred, s_lab) + ) + teacher = with_rest( + chain(teacher_logp, fx["teacher_ids"], b, t_pred, t_lab) + ) + total = total + torch.nn.functional.kl_div( + student, teacher, reduction="sum", log_target=True + ) + beta = knobs["csd_mismatch_loss_beta"] + return beta * total * temperature**2 / fx["global_valid_chunks"] + + +def _realized_rows_payload(fx, k): + """Sparse teacher rows holding each position's realized labels plus fillers.""" + teacher_logits = fx["teacher_logits"] + batch, seq_len, vocab = teacher_logits.shape + required = [[set() for _ in range(seq_len)] for _ in range(batch)] + for b in range(batch): + for c in range(int(fx["num_chunks"][b])): + ss = int(fx["s_spans"][b, c, 0]) + ts, te = fx["t_spans"][b, c].tolist() + shift = fx["knobs"]["kl_chunk_shift"] and ss > 0 and ts > 0 + for label in range(ts, te): + pred = label - 1 if shift else label + required[b][pred].add(int(fx["teacher_ids"][b, label])) + ids = torch.empty(batch, seq_len, k, dtype=torch.long) + for b in range(batch): + for pos in range(seq_len): + row = set(required[b][pos]) + filler = (t for t in range(vocab) if t not in row) + while len(row) < k: + row.add(next(filler)) + ids[b, pos] = torch.tensor(sorted(row)) + temperature = fx["knobs"]["temperature"] + return ( + teacher_logits.gather(-1, ids), + ids.to(torch.int32), + torch.logsumexp(teacher_logits / temperature, dim=-1), + None, + ) + + +def test_csd_realized_only_mismatch_matches_chain_reference(tmp_path): + """Realized-only mismatch loss and gradient match an independent chain KL.""" + fx = _build_case("mismatch_only", str(tmp_path)) + fx["knobs"]["temperature"] = 1.7 + loss, metrics, grad = _run_csd(fx, cfg_overrides=_REALIZED_ONLY) + + student = fx["student_logits"].clone().requires_grad_(True) + reference = _realized_chain_reference(fx, student) + reference.backward() + torch.testing.assert_close(loss, reference, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(grad, student.grad, rtol=1e-5, atol=1e-6) + assert metrics["num_common_chunks"] == 0 + assert metrics["num_mismatch_chunks"] == 4 + + +def test_csd_realized_only_mismatch_sparse_rows_match_dense(tmp_path): + """Sparse rows holding only the realized labels reproduce the dense loss.""" + fx = _build_case("mismatch_only", str(tmp_path)) + fx["knobs"]["temperature"] = 1.7 + dense_loss, _, dense_grad = _run_csd(fx, cfg_overrides=_REALIZED_ONLY) + sparse_loss, sparse_metrics, sparse_grad = _run_csd( + fx, + teacher_sparse_payload=_realized_rows_payload(fx, k=3), + cfg_overrides={**_REALIZED_ONLY, "teacher_topk_ipc_k": 3}, + ) + torch.testing.assert_close(sparse_loss, dense_loss, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(sparse_grad, dense_grad, rtol=1e-5, atol=1e-6) + assert sparse_metrics["num_mismatch_chunks"] == 4 + + +def test_csd_realized_only_mismatch_leaves_common_chunks_unchanged(tmp_path): + """The option changes mismatch scoring only; 1-to-1 chunks are untouched.""" + fx = _build_case("with_mismatch", str(tmp_path)) + prefix_overrides = {**_REALIZED_ONLY, "csd_realized_only_mismatch": False} + _, prefix_metrics, _ = _run_csd(fx, cfg_overrides=prefix_overrides) + _, realized_metrics, _ = _run_csd(fx, cfg_overrides=_REALIZED_ONLY) + assert realized_metrics["num_common_chunks"] == prefix_metrics["num_common_chunks"] + assert realized_metrics["kl_common_per_chunk"] == pytest.approx( + prefix_metrics["kl_common_per_chunk"] + ) + assert ( + realized_metrics["kl_partition_last_per_chunk"] + != prefix_metrics["kl_partition_last_per_chunk"] + ) + + +@pytest.mark.parametrize( + "removed_key", + [ + "csd_position_0_kl", + "csd_mismatch_pos0_alpha", + "csd_mismatch_pos0_weight", + "csd_mismatch_loss_scale", + ], +) +def test_csd_rejects_removed_position_0_keys(removed_key): + """Configs written for the removed position-0 term fail instead of being ignored.""" + from nemo_rl.algorithms.loss.csd import CSDLossConfig + + with pytest.raises(ValueError, match=f"loss_fn.{removed_key} was removed"): + CSDLossConfig.model_validate({removed_key: 0.0}) diff --git a/tests/unit/algorithms/csd/test_csd_owner_local.py b/tests/unit/algorithms/csd/test_csd_owner_local.py new file mode 100644 index 00000000000..66361f7d454 --- /dev/null +++ b/tests/unit/algorithms/csd/test_csd_owner_local.py @@ -0,0 +1,798 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Owner-local CSD evaluation against the replicated reference. + +The student logits stay in their load-balanced CP shard and TP vocabulary +shard. Summed over CP ranks, the owner-local loss and gradient must equal the +replicated TP1/CP1 loss with no ``1/CP`` factor, at every TP/CP size. +""" + +import os +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from nemo_rl.algorithms.loss.csd import ( + CSDLossConfig, + CSDLossFn, + owner_local_config_violations, +) +from nemo_rl.algorithms.x_token import loss_utils +from nemo_rl.algorithms.x_token.loss_utils import ( + LocalizedAlignment, + SparseTeacherRowReader, + _chunk_ids_to_spans, + prepare_csd_owner_local_loss_input, +) +from nemo_rl.algorithms.xtoken_off_policy_distillation import ( + _build_teacher_force_include_token_id_slots, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.model_utils import ( + _get_tokens_on_this_cp_rank, + cp_load_balanced_segments, + distributed_vocab_topk_logz_force, +) +from tests.unit.algorithms.csd.test_csd_loss_parity import _write_subtok_tables + +V = 16 # student and teacher vocabulary (ids 0..9 are 1-to-1 common) +SEQ = 8 +TEMPERATURE = 1.3 +SPARSE_K = 4 + +# (student_start, student_end, teacher_start, teacher_end) per chunk: +# c0 common at position 0 (unshifted) and c1 1-to-2 shifted share teacher +# row 0, which therefore needs two forced labels; c2 2-to-1; c3 common; +# c4 2-to-2 whose student prefix row 5 and final row 6 sit on different CP +# ranks at CP=2. +CHUNKS = [(0, 1, 0, 1), (1, 2, 1, 3), (3, 5, 4, 5), (5, 6, 5, 6), (6, 8, 6, 8)] +STUDENT_IDS = [[0, 5, 0, 9, 5, 2, 3, 1], [1, 5, 0, 9, 6, 4, 7, 2]] +TEACHER_IDS = [[0, 8, 3, 0, 5, 2, 3, 1], [1, 8, 4, 0, 6, 4, 7, 2]] + + +def _fixture(tmpdir: str, *, seed: int = 11) -> dict: + os.makedirs(tmpdir, exist_ok=True) + torch.manual_seed(seed) + fwd_path, rev_path = _write_subtok_tables(tmpdir, V, V) + student_chunk_id = torch.full((2, SEQ), -1, dtype=torch.long) + teacher_chunk_id = torch.full((2, SEQ), -1, dtype=torch.long) + for k, (ss, se, ts, te) in enumerate(CHUNKS): + student_chunk_id[:, ss:se] = k + teacher_chunk_id[:, ts:te] = k + return { + "fwd_path": fwd_path, + "rev_path": rev_path, + "student_logits": torch.randn(2, SEQ, V), + "teacher_logits": torch.randn(2, SEQ, V) * 2.0, + # A same-tokenizer teacher for the multi-teacher case. + "same_vocab_teacher_logits": torch.randn(2, SEQ, V), + "student_ids": torch.tensor(STUDENT_IDS), + "teacher_ids": torch.tensor(TEACHER_IDS), + "student_chunk_id": student_chunk_id, + "teacher_chunk_id": teacher_chunk_id, + "pair_valid": torch.ones((2, len(CHUNKS)), dtype=torch.bool), + "sample_mask": torch.ones(2, dtype=torch.long), + } + + +def _loss_config( + fx: dict, + *, + owner_local: bool, + k: int, + realized_only: bool = False, + noise_filter_topk: int = 0, + dynamic_loss_scaling: bool = True, + same_vocab_teacher: bool = False, +) -> dict: + num_teachers = 2 if same_vocab_teacher else 1 + return { + "temperature": TEMPERATURE, + "vocab_topk": 8, + "reverse_kl": False, + "kl_loss_weight": 1.0, + "ce_loss_scale": 0.3, + "dynamic_loss_scaling": dynamic_loss_scaling, + "kd_loss_mode": "sum", + "student_vocab_size": V, + "teacher_vocab_sizes": [V] * num_teachers, + "projection_matrix_paths": [None] * num_teachers, + "teacher_weights": [1.0, 0.5][:num_teachers], + "common_indices_from_subtoks": True, + "student_to_teacher_table_paths": [fx["fwd_path"], None][:num_teachers], + "teacher_to_student_table_paths": [fx["rev_path"], None][:num_teachers], + "kl_chunk_shift": True, + "csd_loss_fn": "jsd", + "csd_last_pos_loss_fn": "kl", + "csd_jsd_beta": 0.4, + "csd_mismatch_loss_beta": 2.0, + "csd_noise_filter_topk": noise_filter_topk, + "csd_realized_only_mismatch": realized_only, + "teacher_topk_ipc_k": k, + "teacher_topk_ipc_keep_realized": True, + # Small tiles so the batching loops run more than once. + "csd_selected_row_chunk_size": 3, + "csd_support_chunk_batch_size": 2, + "owner_local": owner_local, + } + + +def _batch(fx: dict) -> BatchedDataDict: + token_mask = torch.ones((2, SEQ), dtype=torch.long) + return BatchedDataDict( + { + "input_ids": fx["student_ids"], + "input_lengths": torch.full((2,), SEQ), + "token_mask": token_mask, + "kd_token_mask": token_mask, + "sample_mask": fx["sample_mask"], + "teacher_0_input_ids": fx["teacher_ids"], + "teacher_0_token_mask": token_mask, + "alignment_0_pair_valid": fx["pair_valid"], + "alignment_0_pair_is_correct": fx["pair_valid"], + "alignment_0_student_chunk_id": fx["student_chunk_id"], + "alignment_0_teacher_chunk_id": fx["teacher_chunk_id"], + } + ) + + +def _normalizers(fx: dict) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + sample_mask = fx["sample_mask"].float() + valid_toks = torch.tensor(float((SEQ - 1) * sample_mask.sum())) + valid_chunks = (fx["pair_valid"].float() * sample_mask.unsqueeze(-1)).sum() + return sample_mask.sum(), valid_toks, valid_chunks + + +def _sparse_records( + fx: dict, *, k: int, noise_filter_topk: int, teacher_cp: int +) -> list[dict]: + """Sparse teacher records as the Megatron sender lays them out.""" + force = _build_teacher_force_include_token_id_slots( + _batch(fx), teacher_idx=0, kl_chunk_shift=True + ) + names = ( + "topk_logits", + "topk_indices", + "log_z", + "natural_tail_indices", + "forced_logits", + "forced_indices", + "forced_in_topk", + ) + full = dict( + zip( + names, + distributed_vocab_topk_logz_force( + fx["teacher_logits"], + force, + k, + None, + vocab_start_index=0, + vocab_end_index=V, + real_vocab_size=V, + temperature=TEMPERATURE, + noise_filter_k=noise_filter_topk, + ), + ) + ) + full["log_z"] = full["log_z"].unsqueeze(-1) + entries: list[dict] = [{"teacher_shards": []} for _ in range(2)] + local_seq = SEQ // teacher_cp + for cp_rank in range(teacher_cp): + local = { + name: _get_tokens_on_this_cp_rank(value, cp_rank, teacher_cp, seq_dim=1) + .unsqueeze(0) + .contiguous() + for name, value in full.items() + } + for b in range(2): + entries[b]["teacher_shards"].append( + { + "transport": "sparse_topk", + **{f"{name}_ipc": value for name, value in local.items()}, + "buf_idx": 0, + "sample_index_in_buf": b, + "topk_shape": (local_seq, k), + "forced_shape": (local_seq, 2), + "scalar_shape": (local_seq,), + "k": k, + "force_width": 2, + "temperature": TEMPERATURE, + "real_vocab_size": V, + "forced_in_topk_k": noise_filter_topk or k, + "full_seq_len": SEQ, + "sequence_segments": cp_load_balanced_segments( + SEQ, cp_rank, teacher_cp + ), + } + ) + return entries + + +def _dense_records(teacher_logits: torch.Tensor, *, tp: int, cp: int) -> list[dict]: + """Rectangular dense teacher records with contiguous CP windows.""" + local_seq, local_vocab = SEQ // cp, V // tp + entries: list[dict] = [{"teacher_shards": []} for _ in range(2)] + for cp_rank in range(cp): + for tp_rank in range(tp): + seq = slice(cp_rank * local_seq, (cp_rank + 1) * local_seq) + vocab = slice(tp_rank * local_vocab, (tp_rank + 1) * local_vocab) + storage = teacher_logits[:, seq, vocab].unsqueeze(0).contiguous() + for b in range(2): + entries[b]["teacher_shards"].append( + { + "payload_ipc": storage, + "buf_idx": 0, + "sample_index_in_buf": b, + "actual_shape": (local_seq, local_vocab), + "tp_rank": tp_rank, + "tp_size": tp, + "cp_rank": cp_rank, + "cp_size": cp, + "vocab_start_index": vocab.start, + "vocab_end_index": vocab.stop, + "global_seq_start": seq.start, + "full_seq_len": SEQ, + "full_vocab_size": V, + } + ) + return entries + + +def _replicated_reference( + fx: dict, *, same_vocab_teacher: bool = False, **cfg_kwargs +) -> tuple[torch.Tensor, torch.Tensor, dict]: + """TP1/CP1 replicated CSD loss, gradient and metrics with a dense teacher.""" + loss_fn = CSDLossFn( + _loss_config( + fx, + owner_local=False, + k=0, + same_vocab_teacher=same_vocab_teacher, + **cfg_kwargs, + ) + ) + batch = _batch(fx) + logits = fx["student_logits"].clone().requires_grad_(True) + max_pairs = int(fx["pair_valid"].shape[1]) + aligns = { + 0: LocalizedAlignment( + sample_mask=fx["sample_mask"], + pair_valid=fx["pair_valid"], + pair_is_correct=fx["pair_valid"], + student_input_ids=fx["student_ids"], + student_token_mask=batch["token_mask"], + student_kd_token_mask=batch["kd_token_mask"], + teacher_input_ids=fx["teacher_ids"], + student_spans=_chunk_ids_to_spans(fx["student_chunk_id"], max_pairs), + teacher_spans=_chunk_ids_to_spans(fx["teacher_chunk_id"], max_pairs), + num_chunks=torch.full((2,), max_pairs), + ) + } + teachers = {0: fx["teacher_logits"]} + if same_vocab_teacher: + aligns[1] = LocalizedAlignment( + sample_mask=fx["sample_mask"], + student_input_ids=fx["student_ids"], + student_token_mask=batch["token_mask"], + student_kd_token_mask=batch["kd_token_mask"], + ) + teachers[1] = fx["same_vocab_teacher_logits"] + valid_seqs, valid_toks, valid_chunks = _normalizers(fx) + loss, metrics = loss_fn( + batch, + valid_seqs, + valid_toks, + logits, + logits, + teachers, + aligns, + global_valid_chunks_by_idx={0: valid_chunks}, + global_valid_kd_toks=valid_toks, + ) + loss.backward() + return loss.detach(), logits.grad.detach(), metrics + + +def _groups(rank: int, tp_size: int, cp_size: int): + """Return ``(tp_rank, cp_rank, tp_group, cp_group)`` on a [TP, CP] grid.""" + tp_rank, cp_rank = rank // cp_size, rank % cp_size + tp_group = cp_group = None + for cp_index in range(cp_size): + group = dist.new_group( + ranks=[tp_index * cp_size + cp_index for tp_index in range(tp_size)] + ) + if cp_index == cp_rank: + tp_group = group + for tp_index in range(tp_size): + group = dist.new_group( + ranks=[tp_index * cp_size + cp_index for cp_index in range(cp_size)] + ) + if tp_index == tp_rank: + cp_group = group + return ( + tp_rank, + cp_rank, + tp_group if tp_size > 1 else None, + cp_group if cp_size > 1 else None, + ) + + +def _owner_local_step( + fx: dict, + *, + transport: str, + tp_rank: int, + cp_rank: int, + tp_size: int, + cp_size: int, + tp_group, + cp_group, + same_vocab_teacher: bool = False, + records: list[dict] | None = None, + **cfg_kwargs, +) -> tuple[torch.Tensor, torch.Tensor, dict, torch.Tensor]: + """One rank's owner-local loss, local gradient and metrics.""" + k = SPARSE_K if transport == "sparse" else 0 + loss_fn = CSDLossFn( + _loss_config( + fx, + owner_local=True, + k=k, + same_vocab_teacher=same_vocab_teacher, + **cfg_kwargs, + ) + ) + positions = _get_tokens_on_this_cp_rank( + torch.arange(SEQ), cp_rank, cp_size, seq_dim=0 + ) + vocab_width = V // tp_size + local_logits = ( + fx["student_logits"][ + :, positions, tp_rank * vocab_width : (tp_rank + 1) * vocab_width + ] + .clone() + .requires_grad_(True) + ) + batch = _batch(fx) + if records is None: + records = ( + _sparse_records( + fx, + k=k, + noise_filter_topk=cfg_kwargs.get("noise_filter_topk", 0), + teacher_cp=2, + ) + if transport == "sparse" + else _dense_records(fx["teacher_logits"], tp=2, cp=2) + ) + batch[f"teacher_0_{transport if transport == 'sparse' else 'full'}_logits_ipc"] = ( + records + ) + if same_vocab_teacher: + batch["teacher_1_full_logits_ipc"] = [{"teacher_shards": []}] * 2 + ( + student_logits_contig, + teacher_full, + row_payloads, + aligns, + tp_g, + cp_g, + dp_cp_group, + ) = prepare_csd_owner_local_loss_input( + local_logits, + batch, + cross_tokenizer_teacher_mask=[True, False][: 2 if same_vocab_teacher else 1], + vocab_parallel_group=tp_group, + context_parallel_group=cp_group, + ) + valid_seqs, valid_toks, valid_chunks = _normalizers(fx) + loss, metrics = loss_fn( + batch, + valid_seqs, + valid_toks, + local_logits, + student_logits_contig, + teacher_full, + aligns, + teacher_row_payloads_by_idx=row_payloads, + global_valid_chunks_by_idx={0: valid_chunks}, + global_valid_kd_toks=valid_toks, + tp_group=tp_g, + cp_group=cp_g, + dp_cp_group=dp_cp_group, + ) + loss.backward() + return loss.detach(), local_logits.grad.detach(), metrics, positions + + +def _reassemble_gradient( + local_grad: torch.Tensor, *, world_size: int, tp_size: int, cp_size: int +) -> torch.Tensor: + gathered = [torch.empty_like(local_grad) for _ in range(world_size)] + dist.all_gather(gathered, local_grad.contiguous()) + full = torch.zeros(2, SEQ, V) + vocab_width = V // tp_size + for rank, shard in enumerate(gathered): + tp_rank, cp_rank = rank // cp_size, rank % cp_size + positions = _get_tokens_on_this_cp_rank( + torch.arange(SEQ), cp_rank, cp_size, seq_dim=0 + ) + full[:, positions, tp_rank * vocab_width : (tp_rank + 1) * vocab_width] = shard + return full + + +def _init_gloo(rank: int, world_size: int, init_file: str) -> None: + torch.set_num_threads(1) + dist.init_process_group( + backend="gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + ) + + +def _run_topology( + rank: int, + world_size: int, + init_file: str, + tmpdir: str, + tp_size: int, + cp_size: int, + transport: str, + case: str, +) -> None: + _init_gloo(rank, world_size, init_file) + try: + tp_rank, cp_rank, tp_group, cp_group = _groups(rank, tp_size, cp_size) + fx = _fixture(os.path.join(tmpdir, f"rank{rank}")) + cfg_kwargs: dict = {} + same_vocab_teacher = False + if case == "realized_only": + cfg_kwargs["realized_only"] = True + elif case == "noise_filter": + cfg_kwargs["noise_filter_topk"] = 6 + elif case == "fixed_weights": + cfg_kwargs["dynamic_loss_scaling"] = False + elif case == "same_vocab_teacher": + same_vocab_teacher = True + same_vocab = fx["same_vocab_teacher_logits"] + + def contiguous_window(_entries, cp_group, device): + del device + size = dist.get_world_size(cp_group) if cp_group is not None else 1 + rank_in_cp = dist.get_rank(cp_group) if cp_group is not None else 0 + width = SEQ // size + return same_vocab[:, rank_in_cp * width : (rank_in_cp + 1) * width], 0 + + loss_utils.rebuild_teacher_full_logits_from_ipc = contiguous_window + elif case == "no_owned_chunks": + # Every chunk's final student predictor row is held by CP rank 0 + # (rows 0, 1, 6, 7 at CP=2); CP rank 1 still holds prefix rows. + fx["pair_valid"][:, [2, 3]] = False + + if transport == "sparse": + # Reference: owner-local at TP1/CP1, then cross-checked against the + # replicated path by test_sparse_owner_local_matches_replicated_sparse. + ref_loss, ref_grad, ref_metrics, _ = _owner_local_step( + fx, + transport="sparse", + tp_rank=0, + cp_rank=0, + tp_size=1, + cp_size=1, + tp_group=None, + cp_group=None, + same_vocab_teacher=same_vocab_teacher, + **cfg_kwargs, + ) + else: + ref_loss, ref_grad, ref_metrics = _replicated_reference( + fx, same_vocab_teacher=same_vocab_teacher, **cfg_kwargs + ) + + original_all_gather = dist.all_gather + + def all_gather_without_cp(tensor_list, tensor, group=None, async_op=False): + if not same_vocab_teacher and cp_group is not None and group is cp_group: + raise AssertionError("owner-local KD gathered across CP") + return original_all_gather( + tensor_list, tensor, group=group, async_op=async_op + ) + + dist.all_gather = all_gather_without_cp + try: + loss, local_grad, metrics, _ = _owner_local_step( + fx, + transport=transport, + tp_rank=tp_rank, + cp_rank=cp_rank, + tp_size=tp_size, + cp_size=cp_size, + tp_group=tp_group, + cp_group=cp_group, + same_vocab_teacher=same_vocab_teacher, + **cfg_kwargs, + ) + finally: + dist.all_gather = original_all_gather + + # The fixture must exercise both chunk kinds on every case. + assert ref_metrics["num_common_chunks_t0"] > 0 + assert ref_metrics["num_mismatch_chunks_t0"] > 0 + cp_loss = loss.clone() + if cp_group is not None: + dist.all_reduce(cp_loss, group=cp_group) + torch.testing.assert_close(cp_loss, ref_loss, rtol=2e-5, atol=2e-6) + for key in ("loss", "kl_loss", "ce_loss", "kl_loss_scale", "accuracy"): + assert metrics[key] == pytest.approx( + ref_metrics[key], rel=2e-5, abs=2e-6 + ), key + for key in ("num_common_chunks_t0", "num_mismatch_chunks_t0"): + assert metrics[key] == ref_metrics[key], key + full_grad = _reassemble_gradient( + local_grad, world_size=world_size, tp_size=tp_size, cp_size=cp_size + ) + torch.testing.assert_close(full_grad, ref_grad, rtol=5e-5, atol=5e-6) + finally: + dist.destroy_process_group() + + +_TOPOLOGIES = [(1, 1), (2, 1), (1, 2), (2, 2)] + + +def _spawn(tmp_path: Path, fn, world_size: int, *args) -> None: + mp.spawn( + fn, + args=( + world_size, + os.fspath(tmp_path / "gloo_init"), + os.fspath(tmp_path), + *args, + ), + nprocs=world_size, + join=True, + ) + + +@pytest.mark.skipif(not dist.is_gloo_available(), reason="gloo backend unavailable") +@pytest.mark.parametrize("transport", ["dense", "sparse"]) +@pytest.mark.parametrize("case", ["prefix_support", "realized_only"]) +@pytest.mark.parametrize("tp_size,cp_size", _TOPOLOGIES, ids=lambda v: str(v)) +def test_owner_local_loss_and_gradient_are_topology_invariant( + tmp_path, tp_size, cp_size, transport, case +): + """CP-summed loss and reassembled gradient equal the reference, no 1/CP.""" + _spawn( + tmp_path, + _run_topology, + tp_size * cp_size, + tp_size, + cp_size, + transport, + case, + ) + + +@pytest.mark.skipif(not dist.is_gloo_available(), reason="gloo backend unavailable") +@pytest.mark.parametrize("transport", ["dense", "sparse"]) +@pytest.mark.parametrize( + "case", ["noise_filter", "fixed_weights", "same_vocab_teacher", "no_owned_chunks"] +) +def test_owner_local_tp2_cp2_variants(tmp_path, transport, case): + _spawn(tmp_path, _run_topology, 4, 2, 2, transport, case) + + +def test_sparse_owner_local_matches_replicated_sparse(tmp_path): + """At TP1/CP1, owner-local sparse equals the replicated sparse path. + + The replicated payload holds one forced label per row, so chunk 0 (whose + teacher row would need two labels) is left out of this comparison. + """ + fx = _fixture(os.fspath(tmp_path)) + fx["pair_valid"][:, 0] = False + records = _sparse_records(fx, k=SPARSE_K, noise_filter_topk=0, teacher_cp=2) + owner_loss, owner_grad, owner_metrics, _ = _owner_local_step( + fx, + transport="sparse", + tp_rank=0, + cp_rank=0, + tp_size=1, + cp_size=1, + tp_group=None, + cp_group=None, + records=records, + ) + + # Materialize the replicated [B, T, K] payload: each row's natural top-k + # with its single forced label kept, exactly as the reader serves it. + force = _build_teacher_force_include_token_id_slots( + _batch(fx), teacher_idx=0, kl_chunk_shift=True + ) + assert int((force[..., 1] >= 0).sum()) == 0 + reader = SparseTeacherRowReader(records, device=torch.device("cpu")) + natural_logits, natural_ids = torch.topk(fx["teacher_logits"], SPARSE_K, dim=-1) + natural_ids, order = natural_ids.sort(dim=-1) + topk_logits = natural_logits.gather(-1, order) + topk_ids = natural_ids.to(torch.int32) + labeled = (force[..., 0] >= 0).nonzero().tolist() + rows = reader.gather_rows( + torch.tensor([b for b, _ in labeled]), + torch.tensor([t for _, t in labeled]), + torch.stack([force[b, t, 0] for b, t in labeled]), + ) + for row, (b, t) in enumerate(labeled): + topk_logits[b, t] = rows[0][row] + topk_ids[b, t] = rows[1][row] + log_z = torch.logsumexp(fx["teacher_logits"] / TEMPERATURE, dim=-1) + + loss_fn = CSDLossFn(_loss_config(fx, owner_local=False, k=SPARSE_K)) + batch = _batch(fx) + logits = fx["student_logits"].clone().requires_grad_(True) + max_pairs = int(fx["pair_valid"].shape[1]) + align = LocalizedAlignment( + sample_mask=fx["sample_mask"], + pair_valid=fx["pair_valid"], + student_input_ids=fx["student_ids"], + student_token_mask=batch["token_mask"], + student_kd_token_mask=batch["kd_token_mask"], + teacher_input_ids=fx["teacher_ids"], + student_spans=_chunk_ids_to_spans(fx["student_chunk_id"], max_pairs), + teacher_spans=_chunk_ids_to_spans(fx["teacher_chunk_id"], max_pairs), + num_chunks=torch.full((2,), max_pairs), + ) + valid_seqs, valid_toks, valid_chunks = _normalizers(fx) + loss, metrics = loss_fn( + batch, + valid_seqs, + valid_toks, + logits, + logits, + {}, + {0: align}, + teacher_sparse_logits_by_idx={0: (topk_logits, topk_ids, log_z, None)}, + global_valid_chunks_by_idx={0: valid_chunks}, + global_valid_kd_toks=valid_toks, + ) + loss.backward() + + torch.testing.assert_close(owner_loss, loss.detach(), rtol=2e-5, atol=2e-6) + torch.testing.assert_close(owner_grad, logits.grad, rtol=5e-5, atol=5e-6) + assert owner_metrics["num_mismatch_chunks_t0"] == metrics["num_mismatch_chunks_t0"] + assert owner_metrics["num_common_chunks_t0"] == metrics["num_common_chunks_t0"] + + +def test_owner_local_ce_matches_replicated_ce(tmp_path): + fx = _fixture(os.fspath(tmp_path)) + fx["sample_mask"] = torch.tensor([1, 0]) + batch = _batch(fx) + _, valid_toks, _ = _normalizers(fx) + owner = CSDLossFn(_loss_config(fx, owner_local=True, k=0)) + replicated = CSDLossFn(_loss_config(fx, owner_local=False, k=0)) + logits = fx["student_logits"].clone().requires_grad_(True) + reference_logits = fx["student_logits"].clone().requires_grad_(True) + + owner_ce = owner._compute_owner_local_ce( + logits, batch, valid_toks, tp_group=None, cp_group=None + ) + reference_ce = replicated._compute_ce(reference_logits, batch, valid_toks) + owner_ce.backward() + reference_ce.backward() + + torch.testing.assert_close(owner_ce, reference_ce, rtol=1e-6, atol=1e-6) + torch.testing.assert_close(logits.grad, reference_logits.grad, rtol=1e-6, atol=1e-7) + assert torch.equal(logits.grad[1], torch.zeros_like(logits.grad[1])) + + +def _run_error_agreement( + rank: int, world_size: int, init_file: str, tmpdir: str +) -> None: + _init_gloo(rank, world_size, init_file) + try: + _, cp_rank, _, cp_group = _groups(rank, 1, world_size) + fx = _fixture(os.path.join(tmpdir, f"rank{rank}")) + records = _sparse_records(fx, k=SPARSE_K, noise_filter_topk=0, teacher_cp=2) + if cp_rank == 1: + records[0]["teacher_shards"][0]["transport"] = "unknown" + try: + _owner_local_step( + fx, + transport="sparse", + tp_rank=0, + cp_rank=cp_rank, + tp_size=1, + cp_size=world_size, + tp_group=None, + cp_group=cp_group, + records=records, + ) + except (RuntimeError, ValueError) as exc: + Path(tmpdir, f"error_{rank}.txt").write_text(f"{type(exc).__name__}: {exc}") + else: + Path(tmpdir, f"error_{rank}.txt").write_text("no error") + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_gloo_available(), reason="gloo backend unavailable") +def test_owner_local_setup_failure_raises_on_every_rank(tmp_path): + """A setup error on one CP rank makes every rank raise instead of hanging.""" + _spawn(tmp_path, _run_error_agreement, 2) + errors = [(tmp_path / f"error_{rank}.txt").read_text() for rank in range(2)] + assert errors[0].startswith("RuntimeError") and "another" in errors[0] + assert errors[1].startswith("ValueError") and "transport" in errors[1] + + +def test_owner_local_config_violations_cover_the_restrictions(): + base = {"student_vocab_size": 8, "teacher_topk_ipc_k": 128} + assert owner_local_config_violations(CSDLossConfig(**base)) == [] + cases = { + "kd_loss_mode": ("averaged_logits", "kd_loss_mode"), + "sum_weights_metric": ("ce", "sum_weights_metric"), + "csd_pure_alm": (True, "csd_pure_alm"), + "csd_loss_fn": ("bce", "csd_loss_fn='bce'"), + "csd_last_pos_loss_fn": ("bce", "csd_last_pos_loss_fn='bce'"), + "teacher_topk_ipc_keep_realized": (False, "keep_realized=false"), + } + for field, (value, message) in cases.items(): + violations = owner_local_config_violations( + CSDLossConfig(**{**base, field: value}) + ) + assert len(violations) == 1 and message in violations[0], field + # BCE on the binary realized-only partitions, keep_realized implied by + # realized-only mismatch, and K=1 are all supported. + assert ( + owner_local_config_violations( + CSDLossConfig( + **base, + csd_last_pos_loss_fn="bce", + csd_realized_only_mismatch=True, + teacher_topk_ipc_keep_realized=False, + ) + ) + == [] + ) + assert ( + owner_local_config_violations( + CSDLossConfig( + student_vocab_size=8, + teacher_topk_ipc_k=1, + teacher_topk_ipc_keep_realized=False, + ) + ) + == [] + ) + with pytest.raises(ValueError, match="does not support csd_pure_alm=true"): + CSDLossFn(CSDLossConfig(**base, csd_pure_alm=True, owner_local=True)) + CSDLossFn(CSDLossConfig(**base, csd_pure_alm=True)) + + +def test_owner_local_requires_the_step_global_chunk_count(tmp_path): + fx = _fixture(os.fspath(tmp_path)) + loss_fn = CSDLossFn(_loss_config(fx, owner_local=True, k=0)) + with pytest.raises(ValueError, match="step-global valid-chunk count"): + loss_fn._owner_local_sum_kd( + fx["student_logits"], + None, + {}, + {}, + {}, + torch.tensor(1.0), + None, + tp_group=None, + cp_group=None, + ) diff --git a/tests/unit/algorithms/csd/test_csd_sharded_student_memory.py b/tests/unit/algorithms/csd/test_csd_sharded_student_memory.py new file mode 100644 index 00000000000..00eb63cfc91 --- /dev/null +++ b/tests/unit/algorithms/csd/test_csd_sharded_student_memory.py @@ -0,0 +1,292 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import tempfile +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from tests.unit.algorithms.csd.test_csd_loss_parity import _build_case + + +def _loss_config(fx): + knobs = fx["knobs"] + return { + "temperature": knobs["temperature"], + "vocab_topk": 8, + "reverse_kl": knobs["reverse_kl"], + "kl_loss_weight": 1.0, + "ce_loss_scale": 1.0, + "dynamic_loss_scaling": False, + "kd_loss_mode": "sum", + "normalize_teacher_by_vocab": False, + "alpha": 1.0, + "sum_weights_metric": None, + "student_vocab_size": fx["v_s"], + "teacher_vocab_sizes": [fx["v_t"]], + "projection_matrix_paths": [None], + "teacher_weights": [1.0], + "common_indices_from_subtoks": knobs["common_indices_from_subtoks"], + "student_to_teacher_table_paths": [knobs["student_to_teacher_table_path"]], + "teacher_to_student_table_paths": [knobs["teacher_to_student_table_path"]], + "kl_chunk_shift": knobs["kl_chunk_shift"], + "csd_loss_fn": knobs["csd_loss_fn"], + "csd_last_pos_loss_fn": knobs["csd_last_pos_loss_fn"], + "csd_jsd_beta": knobs["csd_jsd_beta"], + "csd_mismatch_loss_beta": knobs["csd_mismatch_loss_beta"], + "csd_noise_filter_topk": knobs["csd_noise_filter_topk"], + "teacher_topk_ipc_k": 0, + "teacher_topk_ipc_support_mode": "row_topk", + "teacher_topk_ipc_keep_realized": False, + } + + +def _run_tp2_cp2_csd_student_memory( + rank: int, + world_size: int, + init_file: str, + fixture_path: str, +): + from nemo_rl.algorithms.loss import csd as loss_functions_module + from nemo_rl.algorithms.loss.csd import ( + CSDLossFn, + ) + from nemo_rl.algorithms.x_token import loss_utils as shared_loss_utils + from nemo_rl.algorithms.x_token.loss_utils import LocalizedAlignment + + torch.set_num_threads(1) + dist.init_process_group( + backend="gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + ) + try: + tp_size = cp_size = 2 + tp_index, cp_index = rank // cp_size, rank % cp_size + tp_group = None + for cp_coordinate in range(cp_size): + group = dist.new_group( + ranks=[ + tp_coordinate * cp_size + cp_coordinate + for tp_coordinate in range(tp_size) + ] + ) + if cp_coordinate == cp_index: + tp_group = group + cp_group = None + for tp_coordinate in range(tp_size): + group = dist.new_group( + ranks=[ + tp_coordinate * cp_size + cp_coordinate + for cp_coordinate in range(cp_size) + ] + ) + if tp_coordinate == tp_index: + cp_group = group + assert tp_group is not None + assert cp_group is not None + + fx = torch.load(fixture_path, weights_only=False) + loss_fn = CSDLossFn(_loss_config(fx)) + batch_size, student_seq_len, student_vocab_size = fx["student_logits"].shape + teacher_seq_len = fx["teacher_logits"].shape[1] + global_valid_chunks = torch.tensor(float(fx["global_valid_chunks"])) + + def alignment(student_ids, teacher_ids): + return LocalizedAlignment( + sample_mask=torch.ones(batch_size, dtype=torch.bool), + pair_valid=fx["pair_valid"], + student_input_ids=student_ids, + teacher_input_ids=teacher_ids, + student_spans=fx["s_spans"], + teacher_spans=fx["t_spans"], + num_chunks=fx["num_chunks"], + ) + + reference_logits = fx["student_logits"].clone().requires_grad_(True) + reference_loss, reference_metrics = loss_fn._compute_csd_loss( + 0, + reference_logits, + fx["teacher_logits"], + alignment(fx["student_ids"], fx["teacher_ids"]), + teacher_vocab_size=fx["v_t"], + tp_group=None, + cp_group=None, + global_valid_chunks=global_valid_chunks, + ) + reference_loss.backward() + reference_grad = reference_logits.grad.detach().clone() + + student_block = student_seq_len // cp_size + teacher_block = teacher_seq_len // cp_size + vocab_block = student_vocab_size // tp_size + student_start = cp_index * student_block + teacher_start = cp_index * teacher_block + vocab_start = tp_index * vocab_block + student_shard = ( + fx["student_logits"][ + :, + student_start : student_start + student_block, + vocab_start : vocab_start + vocab_block, + ] + .clone() + .requires_grad_(True) + ) + teacher_shard = fx["teacher_logits"][ + :, teacher_start : teacher_start + teacher_block, : + ].clone() + + original_cp_gather = loss_functions_module.allgather_cp_contiguous_tensor + original_full_softmax = shared_loss_utils.vocab_parallel_full_log_softmax + had_vocab_gather = hasattr( + loss_functions_module, + "vocab_parallel_gather_logits", + ) + original_vocab_gather = getattr( + loss_functions_module, + "vocab_parallel_gather_logits", + None, + ) + scalar_gathers = 0 + + def reject_rank3_cp_gather(tensor, group, seq_dim=1): + nonlocal scalar_gathers + if tensor.ndim == 3: + raise AssertionError("CSD attempted a rank-3 CP gather") + scalar_gathers += 1 + return original_cp_gather(tensor, group, seq_dim) + + def reject_vocab_gather(*_args, **_kwargs): + raise AssertionError("CSD attempted a rank-3 TP vocabulary gather") + + loss_functions_module.allgather_cp_contiguous_tensor = reject_rank3_cp_gather + shared_loss_utils.vocab_parallel_full_log_softmax = reject_vocab_gather + loss_functions_module.vocab_parallel_gather_logits = reject_vocab_gather + try: + sharded_loss, sharded_metrics = loss_fn._compute_csd_loss( + 0, + student_shard, + teacher_shard, + alignment( + fx["student_ids"][ + :, student_start : student_start + student_block + ].contiguous(), + fx["teacher_ids"][ + :, teacher_start : teacher_start + teacher_block + ].contiguous(), + ), + teacher_vocab_size=fx["v_t"], + tp_group=tp_group, + cp_group=cp_group, + global_valid_chunks=global_valid_chunks, + ) + finally: + loss_functions_module.allgather_cp_contiguous_tensor = original_cp_gather + shared_loss_utils.vocab_parallel_full_log_softmax = original_full_softmax + if had_vocab_gather: + loss_functions_module.vocab_parallel_gather_logits = ( + original_vocab_gather + ) + else: + delattr(loss_functions_module, "vocab_parallel_gather_logits") + + assert scalar_gathers > 0 + torch.testing.assert_close( + sharded_loss, + reference_loss, + rtol=1.0e-4, + atol=1.0e-4, + ) + (sharded_loss / float(cp_size)).backward() + + gathered_gradients = [torch.empty_like(student_shard.grad) for _ in range(4)] + dist.all_gather(gathered_gradients, student_shard.grad.contiguous()) + reconstructed_grad = torch.zeros_like(reference_grad) + for shard_rank, shard_gradient in enumerate(gathered_gradients): + shard_tp, shard_cp = shard_rank // cp_size, shard_rank % cp_size + reconstructed_grad[ + :, + shard_cp * student_block : (shard_cp + 1) * student_block, + shard_tp * vocab_block : (shard_tp + 1) * vocab_block, + ] = shard_gradient + torch.testing.assert_close( + reconstructed_grad, + reference_grad, + rtol=1.0e-4, + atol=1.0e-4, + ) + + for metric_name in ( + "kl_common_per_chunk", + "kl_partition_last_per_chunk", + "kl_mismatch_scaled_per_chunk", + "top1_acc_per_chunk", + ): + assert sharded_metrics[metric_name] == pytest.approx( + reference_metrics[metric_name], rel=1.0e-4, abs=1.0e-4 + ) + assert sharded_metrics["num_common_chunks"] == 66 + assert sharded_metrics["num_mismatch_chunks"] == 20 + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="gloo backend unavailable", +) +def test_csd_tp2_cp2_is_exact_without_rank3_student_gathers(tmp_path): + """Qualification CSD knobs preserve loss/grad/metrics across TP2/CP2.""" + fx = _build_case("with_mismatch", os.fspath(tmp_path)) + chunks = [(0, 1, 0, 1)] * 33 + [(1, 2, 1, 3), (3, 5, 4, 5)] * 5 + batch_size = fx["student_logits"].shape[0] + fx["s_spans"] = torch.zeros(batch_size, len(chunks), 2, dtype=torch.long) + fx["t_spans"] = torch.zeros(batch_size, len(chunks), 2, dtype=torch.long) + fx["pair_valid"] = torch.ones(batch_size, len(chunks), dtype=torch.bool) + fx["num_chunks"] = torch.full((batch_size,), len(chunks), dtype=torch.long) + for batch_index in range(batch_size): + for chunk_index, ( + student_start, + student_end, + teacher_start, + teacher_end, + ) in enumerate(chunks): + fx["s_spans"][batch_index, chunk_index] = torch.tensor( + [student_start, student_end] + ) + fx["t_spans"][batch_index, chunk_index] = torch.tensor( + [teacher_start, teacher_end] + ) + fx["global_valid_chunks"] = float(batch_size * len(chunks)) + + fixture_path = tmp_path / "fixture.pt" + torch.save(fx, fixture_path) + init_file = tmp_path / "gloo_init" + mp.spawn( + _run_tp2_cp2_csd_student_memory, + args=(4, os.fspath(init_file), os.fspath(fixture_path)), + nprocs=4, + join=True, + ) + + +if __name__ == "__main__": + with tempfile.TemporaryDirectory() as tmpdir: + test_csd_tp2_cp2_is_exact_without_rank3_student_gathers(Path(tmpdir)) + print("CSD TP2/CP2 MEMORY-BOUNDED PARITY PASSED") diff --git a/tests/unit/algorithms/csd/test_distillation.py b/tests/unit/algorithms/csd/test_distillation.py new file mode 100644 index 00000000000..21c1644490c --- /dev/null +++ b/tests/unit/algorithms/csd/test_distillation.py @@ -0,0 +1,2659 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CSD tests for the shared cross-tokenizer training controller. + +Mirrors the style of ``tests/unit/algorithms/test_distillation.py``: +a single ``mock_csd_components`` fixture builds all the Ray/policy/ +data plumbing as ``MagicMock``s, then top-level ``def test_*`` +functions exercise setup, transport, aggregation, and training invariants. +CPU-only, no Ray, no CUDA. + +Cross-tokenizer teachers may ship either dense full-vocab logits or row-wise +top-k logits plus full-vocab logZ over CUDA IPC. Same-vocab teachers remain +dense because their direct-KL and teacher-scoring paths require full logits. +""" + +from __future__ import annotations + +import math +import os +from copy import deepcopy +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from pydantic import ValidationError +from torchdata.stateful_dataloader import StatefulDataLoader + +import nemo_rl.algorithms.xtoken_off_policy_distillation as csd_mod +from nemo_rl.algorithms.loss.csd import CSDLossConfig, CSDLossFn +from nemo_rl.algorithms.xtoken_off_policy_distillation import ( + CSDMasterConfig, + TeacherAlignerConfig, + TeacherConfig, + _build_teacher_force_include_token_id_slots, + _build_teacher_force_include_token_ids, + _default_off_policy_distillation_save_state, + _packing_batch_uids_from_state, + build_csd_lockstep_packing_plan, + build_csd_logical_batch_digest_record, + export_teacher_logits_and_pack, + log_csd_logical_batch_digest, + log_csd_packing_telemetry, + reduce_mb_metric, + setup, + validate, + validate_csd_packing_setup, + validate_csd_tokenizer_reuse, + xtoken_non_student_seq_keys, + xtoken_off_policy_distillation_train, +) +from nemo_rl.data.cross_tokenizer.cross_tokenizer_collate import ( + CrossTokenizerCollatorConfig, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.virtual_cluster import ClusterConfig +from nemo_rl.utils.logger import LoggerConfig + + +def has_gloo() -> bool: + """Whether torch.distributed has a usable gloo backend.""" + return torch.distributed.is_available() and torch.distributed.is_gloo_available() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_batch( + batch_size: int = 1, + t_student: int = 4, + t_teacher: int = 4, + num_teachers: int = 1, +) -> dict: + """Synthetic batch with every teacher-indexed key the packer reads. + + Each teacher ``i`` is cross-tokenizer here, so it carries + ``teacher_{i}_*`` tokenization + ``alignment_{i}_*`` keys (matching what + ``CrossTokenizerCollator`` emits for a cross-tokenizer teacher). + """ + batch = { + "input_ids": torch.zeros((batch_size, t_student), dtype=torch.long), + "input_lengths": torch.full((batch_size,), t_student, dtype=torch.long), + "token_mask": torch.ones((batch_size, t_student), dtype=torch.long), + "sample_mask": torch.ones((batch_size,), dtype=torch.long), + "sample_id": [f"sample-{i}" for i in range(batch_size)], + } + for i in range(num_teachers): + batch[f"teacher_{i}_input_ids"] = torch.zeros( + (batch_size, t_teacher), dtype=torch.long + ) + batch[f"teacher_{i}_input_lengths"] = torch.full( + (batch_size,), t_teacher, dtype=torch.long + ) + batch[f"teacher_{i}_token_mask"] = torch.ones( + (batch_size, t_teacher), dtype=torch.long + ) + batch[f"alignment_{i}_pair_valid"] = torch.ones( + (batch_size, 2), dtype=torch.bool + ) + batch[f"alignment_{i}_pair_is_correct"] = torch.ones( + (batch_size, 2), dtype=torch.bool + ) + batch[f"alignment_{i}_student_exact_partition_mask"] = torch.zeros( + (batch_size, t_student), dtype=torch.bool + ) + batch[f"alignment_{i}_teacher_exact_partition_mask"] = torch.zeros( + (batch_size, t_teacher), dtype=torch.bool + ) + batch[f"alignment_{i}_student_chunk_id"] = torch.zeros( + (batch_size, t_student), dtype=torch.long + ) + batch[f"alignment_{i}_teacher_chunk_id"] = torch.zeros( + (batch_size, t_teacher), dtype=torch.long + ) + batch[f"alignment_{i}_num_chunks"] = torch.tensor( + [1] * batch_size, dtype=torch.long + ) + # validate() pads ragged val batches via BatchedDataDict.size, so the mock + # batches must be BatchedDataDict (the train path reads them as a dict too). + return BatchedDataDict(batch) + + +def _mock_dataloader(num_batches: int) -> MagicMock: + batch = _make_batch() + dl = MagicMock(spec=StatefulDataLoader) + dl.__iter__ = lambda self: iter([batch] * num_batches) + dl.__len__ = MagicMock(return_value=num_batches) + dl.state_dict = MagicMock(return_value={}) + return dl + + +def _make_master_config( + *, + max_num_steps: int = 5, + max_num_epochs: int = 10, + val_period: int = 100, + val_at_start: bool = False, + val_at_end: bool = False, + save_enabled: bool = False, +) -> CSDMasterConfig: + """CSDMasterConfig that passes the setup() backend asserts and the + train loop's lookups. Built via ``CSDMasterConfig.model_construct`` to + bypass strict TypedDict field validation (matches the pattern used + in tests/unit/algorithms/test_rm.py). + """ + return CSDMasterConfig.model_construct( + **{ + "distillation": { + "num_prompts_per_step": 1, + "max_num_steps": max_num_steps, + "max_num_epochs": max_num_epochs, + "seed": 42, + "val_period": val_period, + "val_at_start": val_at_start, + "val_at_end": val_at_end, + }, + "policy": { + "automodel_cfg": { + "enabled": True, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + }, + "max_total_sequence_length": 64, + "make_sequence_length_divisible_by": 8, + "train_global_batch_size": 1, + "train_micro_batch_size": 1, + "tokenizer": {"name": "student-tok"}, + }, + "teachers": [ + TeacherConfig( + **{ + "aligner": { + "projection_matrix_path": None, + "student_to_teacher_table_path": "/tmp/forward.pt", + "teacher_to_student_table_path": "/tmp/reverse.pt", + }, + "weight": 1.0, + "automodel_cfg": { + "enabled": True, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + }, + "max_total_sequence_length": 64, + "make_sequence_length_divisible_by": 8, + "train_global_batch_size": 1, + "train_micro_batch_size": 1, + "tokenizer": {"name": "teacher-tok"}, + } + ) + ], + "loss_fn": CSDLossConfig( + **{ + "temperature": 1.0, + "vocab_topk": 8, + "reverse_kl": False, + "kl_loss_weight": 1.0, + "ce_loss_scale": 1.0, + "dynamic_loss_scaling": False, + "kd_loss_mode": "sum", + "sum_weights_metric": None, + "alpha": 1.0, + "normalize_teacher_by_vocab": False, + "kl_chunk_shift": False, + "csd_noise_filter_topk": 0, + "teacher_topk_ipc_k": 0, + "teacher_topk_ipc_support_mode": "row_topk", + "teacher_topk_ipc_keep_realized": True, + } + ), + "data": { + "shuffle": False, + "num_workers": 0, + }, + "logger": LoggerConfig(log_dir="/tmp/logger"), + "cluster": ClusterConfig(num_nodes=1, gpus_per_node=1), + "checkpointing": { + "enabled": save_enabled, + "checkpoint_must_save_by": None, + "save_period": 100, + "metric_name": None, + }, + } + ) + + +def _make_tokenizer(vocab_size: int) -> MagicMock: + tok = MagicMock() + tok.__len__ = MagicMock(return_value=vocab_size) + tok.get_vocab.return_value = { + f"token-{token_id}": token_id for token_id in range(vocab_size) + } + tok.is_fast = True + tok.backend_tokenizer.to_str.return_value = "stable-fast-backend" + tok.special_tokens_map_extended = {} + tok.all_special_ids = [] + tok.chat_template = None + return tok + + +def _enable_lockstep_packing( + cfg: CSDMasterConfig, *, global_batch_size: int = 1, capacity: int = 64 +) -> None: + cfg.distillation["num_prompts_per_step"] = global_batch_size + cfg.data["train"] = { + "dataset_name": "arrow_text", + "characters_per_sample": None, + } + cfg.collator.num_packed_rows = 1 + + def enable_policy_config(policy_config): + policy_config["model_name"] = "transformer-test-model" + policy_config["train_global_batch_size"] = global_batch_size + policy_config["train_micro_batch_size"] = 1 + policy_config["dynamic_batching"] = {"enabled": False} + policy_config["sequence_packing"] = { + "enabled": True, + "train_mb_tokens": capacity, + "logprob_mb_tokens": capacity, + "algorithm": "lockstep_first_fit_decreasing", + "fuse_loss": False, + } + + enable_policy_config(cfg.policy) + for teacher_idx, teacher in enumerate(cfg.teachers): + teacher_config = teacher.policy_config() + enable_policy_config(teacher_config) + cfg.teachers[teacher_idx] = TeacherConfig( + **teacher_config, + aligner=teacher.aligner.model_dump(), + weight=teacher.weight, + ) + + +# --------------------------------------------------------------------------- +# Fixture +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_csd_components(): + student_policy = MagicMock() + student_policy.data_parallel_size = 1 + student_policy.train.return_value = { + "loss": torch.tensor(0.5), + "grad_norm": torch.tensor(1.0), + "all_mb_metrics": { + "global_valid_toks": [10], + "kl_loss": [0.3], + "ipc_reconstruction_fallbacks": [1], + "ipc_reconstruction_fallbacks_t0": [1], + }, + } + + teacher_policy = MagicMock() + teacher_policy.data_parallel_size = 1 + teacher_policy.get_full_logits_ipc.return_value = [ + {"teacher_shards": [{"actual_shape": (4, 32), "dtype": torch.float32}]} + ] + + train_dataloader = _mock_dataloader(num_batches=10) + val_dataloader = _mock_dataloader(num_batches=2) + + # The trainer reads per-teacher metadata off the loss fn to build the + # skip-keys set and drive the teacher-forward loop. One cross-tokenizer + # teacher: non-null projection path (every teacher ships full logits). + loss_fn = MagicMock() + loss_fn.num_teachers = 1 + loss_fn.projection_matrix_paths = ["/tmp/dummy-projection.pt"] + loss_fn.teacher_vocab_sizes = [24] + loss_fn.cfg = CSDLossConfig( + **{ + "temperature": 1.0, + "kl_chunk_shift": False, + "csd_noise_filter_topk": 0, + "teacher_topk_ipc_k": 0, + "teacher_topk_ipc_support_mode": "row_topk", + "teacher_topk_ipc_keep_realized": True, + } + ) + logger = MagicMock() + + checkpointer = MagicMock() + checkpointer.save_optimizer = False + + return SimpleNamespace( + student_policy=student_policy, + teacher_policy=teacher_policy, + train_dataloader=train_dataloader, + val_dataloader=val_dataloader, + loss_fn=loss_fn, + logger=logger, + checkpointer=checkpointer, + save_state=_default_off_policy_distillation_save_state(), + master_config=_make_master_config(), + ) + + +@pytest.fixture(autouse=True) +def _resolve_synthetic_packing_model_metadata(monkeypatch): + """Keep synthetic lockstep configs offline while exercising metadata checks.""" + original = csd_mod._load_csd_packed_model_config + + def load(config): + if config.get("model_name") == "transformer-test-model": + return SimpleNamespace( + model_type="llama", + architectures=["LlamaForCausalLM"], + to_dict=lambda: { + "model_type": "llama", + "architectures": ["LlamaForCausalLM"], + }, + ) + return original(config) + + monkeypatch.setattr(csd_mod, "_load_csd_packed_model_config", load) + + +# --------------------------------------------------------------------------- +# setup() backend & vocab-injection asserts +# --------------------------------------------------------------------------- + + +def _patched_setup_call( + master_config, *, student_vocab=32, teacher_vocab=24, train_batches=4 +): + """Drive setup() with every heavy collaborator patched out.""" + student_tok = _make_tokenizer(student_vocab) + teacher_tok = _make_tokenizer(teacher_vocab) + train_ds = MagicMock() + train_ds.__len__ = MagicMock(return_value=4) + val_ds = MagicMock() + val_ds.__len__ = MagicMock(return_value=2) + + with ( + patch.object(csd_mod, "RayVirtualCluster") as mock_cluster, + patch.object(csd_mod, "Policy") as mock_policy_cls, + patch.object(csd_mod, "Logger"), + patch.object(csd_mod, "CheckpointManager") as mock_cp_cls, + patch.object(csd_mod, "TokenAligner"), + patch.object(csd_mod, "CSDCrossTokenizerCollator") as mock_collator_cls, + patch.object(csd_mod, "CSDLossFn") as mock_loss_cls, + patch.object(csd_mod, "StatefulDataLoader") as mock_dl_cls, + patch.object(csd_mod, "assert_teacher_student_batch_grid"), + patch.object(csd_mod, "assert_xtoken_ipc_node_local"), + ): + mock_cp_cls.return_value.get_latest_checkpoint_path.return_value = None + mock_cp_cls.return_value.load_training_info.return_value = None + mock_cp_cls.return_value.get_resume_paths.return_value = (None, None) + train_dl = MagicMock(spec=StatefulDataLoader) + train_dl.__len__ = MagicMock(return_value=train_batches) + val_dl = MagicMock(spec=StatefulDataLoader) + val_dl.__len__ = MagicMock(return_value=2) + mock_dl_cls.side_effect = [train_dl, val_dl] + mock_policy_cls.side_effect = lambda *a, **kw: MagicMock(data_parallel_size=1) + + result = setup( + master_config, + student_tokenizer=student_tok, + teacher_tokenizers=[teacher_tok], + train_dataset=train_ds, + val_dataset=val_ds, + ) + return result, { + "cluster": mock_cluster, + "policy": mock_policy_cls, + "loss": mock_loss_cls, + "collator": mock_collator_cls, + "checkpointer": mock_cp_cls, + } + + +def test_teacher_aligner_config_defaults(): + teacher = TeacherConfig(model_name="teacher") + + assert isinstance(teacher.aligner, TeacherAlignerConfig) + assert teacher.aligner.projection_matrix_path is None + assert teacher.aligner.drop_first_assistant_chunk_kl is False + assert teacher.aligner.student_to_teacher_table_path is None + assert teacher.aligner.teacher_to_student_table_path is None + + +def test_teacher_aligner_config_explicit_values_serialize_and_stay_out_of_policy(): + teacher = TeacherConfig( + model_name="teacher", + aligner={ + "projection_matrix_path": "/tmp/projection.pt", + "drop_first_assistant_chunk_kl": True, + "student_to_teacher_table_path": "/tmp/forward.pt", + "teacher_to_student_table_path": "/tmp/reverse.pt", + }, + ) + + dumped = teacher.model_dump() + assert dumped["aligner"] == { + "projection_matrix_path": "/tmp/projection.pt", + "drop_first_assistant_chunk_kl": True, + "student_to_teacher_table_path": "/tmp/forward.pt", + "teacher_to_student_table_path": "/tmp/reverse.pt", + } + assert "projection_matrix_path" not in dumped + assert "aligner" not in teacher.policy_config() + + +@pytest.mark.parametrize( + "field", ["student_to_teacher_table_path", "teacher_to_student_table_path"] +) +def test_root_student_to_teacher_table_paths_are_rejected(field): + with pytest.raises( + ValidationError, + match="must be nested under teachers\\[i\\]\\.aligner", + ): + TeacherConfig(**{field: "/tmp/table.pt"}) + + +def test_empty_teachers_list_rejected_at_config_load(): + """An empty ``teachers:`` list fails validation at config load.""" + with pytest.raises(ValidationError) as exc_info: + CSDMasterConfig(teachers=[]) + assert any( + e["loc"] == ("teachers",) and e["type"] == "too_short" + for e in exc_info.value.errors() + ) + + +def test_setup_requires_supported_student_backend(): + cfg = _make_master_config() + cfg.policy["automodel_cfg"]["enabled"] = False + with ( + patch.object(csd_mod, "RayVirtualCluster") as mock_cluster, + pytest.raises(AssertionError), + ): + setup( + cfg, + student_tokenizer=_make_tokenizer(32), + teacher_tokenizers=[_make_tokenizer(24)], + train_dataset=MagicMock(), + val_dataset=None, + ) + assert mock_cluster.call_count == 0 + + +def test_setup_requires_supported_teacher_backend(): + cfg = _make_master_config() + cfg.teachers[0].automodel_cfg["enabled"] = False + with ( + patch.object(csd_mod, "RayVirtualCluster") as mock_cluster, + pytest.raises(AssertionError), + ): + setup( + cfg, + student_tokenizer=_make_tokenizer(32), + teacher_tokenizers=[_make_tokenizer(24)], + train_dataset=MagicMock(), + val_dataset=None, + ) + assert mock_cluster.call_count == 0 + + +def test_setup_router_replay_requires_megatron_student(): + cfg = _make_master_config() + cfg.policy["router_replay"] = {"enabled": True} + with ( + patch.object(csd_mod, "RayVirtualCluster") as mock_cluster, + pytest.raises(ValueError, match="requires the Megatron student policy backend"), + ): + setup( + cfg, + student_tokenizer=_make_tokenizer(32), + teacher_tokenizers=[_make_tokenizer(24)], + train_dataset=MagicMock(), + val_dataset=None, + ) + assert mock_cluster.call_count == 0 + + +def _automodel_cfg(*, enabled=True, tp=1, cp=1): + return { + "enabled": enabled, + "tensor_parallel_size": tp, + "context_parallel_size": cp, + } + + +def _megatron_cfg(*, enabled=True, tp=1, pp=1, cp=1): + return { + "enabled": enabled, + "tensor_model_parallel_size": tp, + "pipeline_model_parallel_size": pp, + "context_parallel_size": cp, + } + + +def test_entity_parallelism_automodel_returns_tp_cp_pp(): + # DTensor has no pipeline axis, so it always reports pp == 1. + cfg = {"automodel_cfg": _automodel_cfg(tp=4, cp=2)} + assert csd_mod._xtoken_entity_parallelism(cfg, label="teachers[0]") == (4, 2, 1) + + +def test_entity_parallelism_megatron_pp1cp1_returns_tp_cp_pp(): + cfg = { + "automodel_cfg": _automodel_cfg(enabled=False), + "megatron_cfg": _megatron_cfg(tp=2, pp=1, cp=1), + } + assert csd_mod._xtoken_entity_parallelism(cfg, label="teachers[0]") == (2, 1, 1) + + +def test_entity_parallelism_megatron_accepts_pp_gt_1(): + # PP is supported: only the last stage holds logits, so only it contributes + # full-logits IPC handles and the earlier stages drop out in + # aggregate_per_sample_handles. + cfg = { + "automodel_cfg": _automodel_cfg(enabled=False), + "megatron_cfg": _megatron_cfg(tp=2, pp=2, cp=1), + } + assert csd_mod._xtoken_entity_parallelism(cfg, label="teachers[0]") == (2, 1, 2) + + +def test_entity_parallelism_megatron_accepts_cp_gt_1(): + # Megatron TP/CP/PP are all supported (the loss is parallelism-invariant). + cfg = { + "automodel_cfg": _automodel_cfg(enabled=False), + "megatron_cfg": _megatron_cfg(tp=2, pp=1, cp=2), + } + assert csd_mod._xtoken_entity_parallelism(cfg, label="teachers[0]") == (2, 2, 1) + + +def test_entity_parallelism_rejects_neither_backend(): + # Neither supported training backend is enabled. + cfg = {"automodel_cfg": _automodel_cfg(enabled=False)} + with pytest.raises(AssertionError, match="either Automodel"): + csd_mod._xtoken_entity_parallelism(cfg, label="teachers[0]") + + +def test_setup_injects_vocab_sizes_into_loss_config(): + cfg = _make_master_config() + cfg.teachers[0].aligner.student_to_teacher_table_path = "/tmp/forward.pt" + cfg.teachers[0].aligner.teacher_to_student_table_path = "/tmp/reverse.pt" + original_loss_cfg = deepcopy(cfg.loss_fn) + + _, mocks = _patched_setup_call(cfg, student_vocab=128, teacher_vocab=256) + + mocks["loss"].assert_called_once() + injected_cfg = mocks["loss"].call_args.args[0] + assert injected_cfg["student_vocab_size"] == 128 + # Per-teacher metadata is injected as parallel lists (one teacher here). + assert injected_cfg["teacher_vocab_sizes"] == [256] + assert injected_cfg["projection_matrix_paths"] == [None] + assert injected_cfg["teacher_weights"] == [1.0] + assert injected_cfg["student_to_teacher_table_paths"] == ["/tmp/forward.pt"] + assert injected_cfg["teacher_to_student_table_paths"] == ["/tmp/reverse.pt"] + assert mocks["collator"].call_args.kwargs[ + "drop_first_assistant_chunk_kl_by_teacher" + ] == [False] + # Original master_config not mutated by the injection. + assert cfg.loss_fn == original_loss_cfg + + +_COLLATOR_OPTION_NAMES = ( + "require_routed_experts", + "reject_overlength_text", + "add_special_tokens_in_text", + "clip_overlength_text", + "ce_covers_assistant_turn", +) + + +@pytest.mark.parametrize( + ("algorithm", "collator", "expected_options"), + [ + pytest.param( + "csd", + {"mode": "text"}, + { + "require_routed_experts": False, + "reject_overlength_text": True, + "add_special_tokens_in_text": True, + "clip_overlength_text": True, + }, + id="csd-text", + ), + pytest.param( + "csd", + {"mode": "chat"}, + {"require_routed_experts": False, "reject_overlength_text": True}, + id="csd-chat", + ), + pytest.param( + "csd", + { + "mode": "chat", + "include_thinking_in_loss": True, + "native_thinking_alignment": True, + }, + {"require_routed_experts": False, "reject_overlength_text": True}, + id="csd-native-thinking-chat", + ), + pytest.param("xtoken", {"mode": "chat"}, {}, id="xtoken-chat"), + ], +) +def test_setup_passes_per_algorithm_collator_options( + algorithm, collator, expected_options +): + cfg = _make_master_config() + cfg.collator = CrossTokenizerCollatorConfig(**collator) + if algorithm == "xtoken": + cfg.algorithm = "xtoken" + cfg.teachers[0].aligner.projection_matrix_path = "/tmp/projection.pt" + + with ( + patch.object(csd_mod, "CrossTokenizerCollator") as shared_collator_cls, + patch.object(csd_mod, "CrossTokenizerDistillationLossFn"), + ): + _, mocks = _patched_setup_call(cfg) + + # CSD builds its subclass; X-Token builds the shared collator. + used, unused = (mocks["collator"], shared_collator_cls) + if algorithm == "xtoken": + used, unused = unused, used + unused.assert_not_called() + used.assert_called_once() + kwargs = used.call_args.kwargs + assert kwargs["config"] == cfg.collator + assert { + name: kwargs[name] for name in _COLLATOR_OPTION_NAMES if name in kwargs + } == expected_options + + +def test_setup_requires_routed_experts_when_router_replay_is_enabled(): + cfg = _make_master_config() + cfg.policy["automodel_cfg"]["enabled"] = False + cfg.policy["megatron_cfg"] = _megatron_cfg() + cfg.policy["router_replay"] = {"enabled": True} + + _, mocks = _patched_setup_call(cfg) + + mocks["collator"].assert_called_once() + assert mocks["collator"].call_args.kwargs["require_routed_experts"] is True + + +def test_setup_sets_derived_train_iters_on_megatron_teacher_and_student(): + cfg = _make_master_config(max_num_steps=10, max_num_epochs=2) + cfg.policy["automodel_cfg"]["enabled"] = False + cfg.policy["megatron_cfg"] = _megatron_cfg() + cfg.teachers[0].automodel_cfg["enabled"] = False + cfg.teachers[0].megatron_cfg = _megatron_cfg() + + _, mocks = _patched_setup_call(cfg, train_batches=3) + + # min(max_num_steps=10, max_num_epochs=2 * train_batches=3) == 6. + teacher_config = mocks["policy"].call_args_list[0].kwargs["config"] + student_config = mocks["policy"].call_args_list[1].kwargs["config"] + assert teacher_config["megatron_cfg"]["train_iters"] == 6 + assert student_config["megatron_cfg"]["train_iters"] == 6 + + +@pytest.mark.parametrize( + "val_dataset, val_period, val_at_start, val_at_end, expect_loader", + [ + ("present", 0, False, False, False), # all gates off + (None, 100, True, True, False), # no dataset + ("present", 100, False, False, True), # val_period > 0 + ("present", 0, True, False, True), # val_at_start + ("present", 0, False, True, True), # val_at_end + ], +) +def test_setup_val_dataloader_gating( + val_dataset, val_period, val_at_start, val_at_end, expect_loader +): + cfg = _make_master_config( + val_period=val_period, val_at_start=val_at_start, val_at_end=val_at_end + ) + + ds = MagicMock() + ds.__len__ = MagicMock(return_value=2) + val_ds = ds if val_dataset is not None else None + student_tok = _make_tokenizer(32) + teacher_tok = _make_tokenizer(24) + train_ds = MagicMock() + train_ds.__len__ = MagicMock(return_value=4) + + with ( + patch.object(csd_mod, "RayVirtualCluster"), + patch.object(csd_mod, "Policy") as mock_policy_cls, + patch.object(csd_mod, "Logger"), + patch.object(csd_mod, "CheckpointManager") as mock_cp_cls, + patch.object(csd_mod, "TokenAligner"), + patch.object(csd_mod, "CSDCrossTokenizerCollator"), + patch.object(csd_mod, "CSDLossFn"), + patch.object(csd_mod, "StatefulDataLoader") as mock_dl_cls, + patch.object(csd_mod, "assert_teacher_student_batch_grid"), + patch.object(csd_mod, "assert_xtoken_ipc_node_local"), + ): + mock_cp_cls.return_value.get_latest_checkpoint_path.return_value = None + mock_cp_cls.return_value.load_training_info.return_value = None + mock_cp_cls.return_value.get_resume_paths.return_value = (None, None) + mock_dl_cls.side_effect = lambda *a, **kw: MagicMock(spec=StatefulDataLoader) + mock_policy_cls.side_effect = lambda *a, **kw: MagicMock(data_parallel_size=1) + + ( + _student, + _teacher, + _train_dl, + val_dl, + *_, + ) = setup( + cfg, + student_tokenizer=student_tok, + teacher_tokenizers=[teacher_tok], + train_dataset=train_ds, + val_dataset=val_ds, + ) + + if expect_loader: + assert val_dl is not None + else: + assert val_dl is None + + +# --------------------------------------------------------------------------- +# Training loop: exit conditions +# --------------------------------------------------------------------------- + + +def _run_train(c): + xtoken_off_policy_distillation_train( + c.student_policy, + [c.teacher_policy], + c.train_dataloader, + c.val_dataloader, + c.loss_fn, + c.logger, + c.checkpointer, + c.save_state, + c.master_config, + ) + + +def test_exit_on_max_steps(mock_csd_components): + mock_csd_components.master_config.distillation["max_num_steps"] = 3 + mock_csd_components.master_config.distillation["max_num_epochs"] = 10 + + _run_train(mock_csd_components) + + assert mock_csd_components.student_policy.train.call_count == 3 + + +def test_train_surfaces_reduced_teacher_routing_metrics(mock_csd_components): + c = mock_csd_components + c.master_config.distillation.update(max_num_steps=1, max_num_epochs=1) + c.val_dataloader = None + c.student_policy.train.return_value["all_mb_metrics"].update( + { + "teacher_0/routed_samples": [2, 3], + "teacher_0/routed_tokens": [7, 11], + "teacher_0/weighted_kl": [0.2, 0.3], + } + ) + + _run_train(c) + + train_log_call = next( + call + for call in c.logger.log_metrics.call_args_list + if call.kwargs.get("prefix") == "train" + ) + metrics = train_log_call.args[0] + assert metrics["teacher_0/routed_samples"] == pytest.approx(5) + assert metrics["teacher_0/routed_tokens"] == pytest.approx(18) + assert metrics["teacher_0/weighted_kl"] == pytest.approx(0.5) + + +def test_ft_save_period_triggers_periodic_saves(mock_csd_components): + """ft_save_period triggers checkpoint saves independent of save_period.""" + c = mock_csd_components + c.master_config.distillation["max_num_steps"] = 5 + c.master_config.distillation["max_num_epochs"] = 1 + c.master_config.checkpointing["enabled"] = True + c.master_config.checkpointing["save_period"] = 100 # only the final step saves + c.master_config.checkpointing["ft_save_period"] = 2 + c.master_config.checkpointing["metric_name"] = None + c.checkpointer.init_tmp_checkpoint.return_value = "/tmp/ft_ckpt_test/tmp_step" + + with patch("nemo_rl.algorithms.xtoken_off_policy_distillation.torch.save"): + _run_train(c) + + # ft_save_period=2 -> steps 2, 4; save_period=100 contributes only the last step (5). + saved_steps = [ + call.args[0] for call in c.checkpointer.init_tmp_checkpoint.call_args_list + ] + assert saved_steps == [2, 4, 5] + + +@pytest.mark.parametrize( + ("offload_student", "offload_optimizer", "restore_method", "unused_method"), + [ + (True, True, "offload_after_refit", "offload_before_refit"), + (False, True, "offload_before_refit", "offload_after_refit"), + ], +) +def test_checkpoint_restores_student_between_teacher_state_after_save( + mock_csd_components, + offload_student, + offload_optimizer, + restore_method, + unused_method, +): + c = mock_csd_components + c.master_config.distillation.update( + max_num_steps=1, + max_num_epochs=1, + offload_student_after_step=offload_student, + ) + c.master_config.policy["offload_optimizer_for_logprob"] = offload_optimizer + c.master_config.checkpointing.update(enabled=True, save_period=1) + c.checkpointer.init_tmp_checkpoint.return_value = "/tmp/residency/tmp_step" + + events = MagicMock() + events.attach_mock(c.student_policy, "student") + events.attach_mock(c.checkpointer, "checkpointer") + with patch("nemo_rl.algorithms.xtoken_off_policy_distillation.torch.save"): + _run_train(c) + + event_names = [event[0] for event in events.mock_calls] + save_index = event_names.index("student.save_checkpoint") + finalize_index = event_names.index("checkpointer.begin_finalization") + restore_indices = [ + index + for index, event_name in enumerate(event_names) + if event_name == f"student.{restore_method}" + ] + + assert save_index < finalize_index < restore_indices[-1] + assert len(restore_indices) == 2 + getattr(c.student_policy, unused_method).assert_not_called() + + +def test_exit_on_max_epochs(mock_csd_components): + # max_num_steps high so it doesn't fire first; max_num_epochs caps the loop. + mock_csd_components.master_config.distillation["max_num_steps"] = 10_000 + mock_csd_components.master_config.distillation["max_num_epochs"] = 2 + # Two batches per epoch * 2 epochs = 4 student.train calls. + mock_csd_components.train_dataloader = _mock_dataloader(num_batches=2) + + _run_train(mock_csd_components) + + assert mock_csd_components.student_policy.train.call_count == 4 + + +def test_exit_on_timeout(mock_csd_components, capsys): + mock_csd_components.master_config.distillation["max_num_steps"] = 100 + + with patch.object(csd_mod, "TimeoutChecker") as mock_timeout_class: + mock_timeout_instance = MagicMock() + # False for 4 steps, then True (timeout). + mock_timeout_instance.check_save.side_effect = [False] * 4 + [True] + mock_timeout_class.return_value = mock_timeout_instance + + _run_train(mock_csd_components) + + # Loop should have run exactly 5 steps before tripping the timeout return. + assert mock_csd_components.student_policy.train.call_count == 5 + + captured = capsys.readouterr() + assert "Timeout reached, stopping training early." in captured.out + + +def test_packing_batch_uid_high_water_mark_survives_multibatch_validation_resume( + mock_csd_components, +): + c = mock_csd_components + _enable_lockstep_packing(c.master_config, global_batch_size=1, capacity=64) + c.master_config.distillation.update( + max_num_steps=1, + max_num_epochs=10, + val_period=1, + val_at_start=False, + val_at_end=False, + ) + c.val_dataloader = _mock_dataloader(num_batches=2) + c.master_config.checkpointing.update( + enabled=True, + save_period=1, + ft_save_period=None, + metric_name=None, + ) + c.checkpointer.init_tmp_checkpoint.return_value = "/tmp/uid_resume/tmp_step" + + with patch("nemo_rl.algorithms.xtoken_off_policy_distillation.torch.save"): + _run_train(c) + + # Reconstruct the state as JSON checkpoint loading would, then continue for + # one more step. The first run consumed one train UID plus two validation + # UIDs, so resume must start at 3 rather than deriving 2 from total_steps. + c.save_state = dict(c.save_state) + c.master_config.distillation["max_num_steps"] = 2 + with patch("nemo_rl.algorithms.xtoken_off_policy_distillation.torch.save"): + _run_train(c) + + observed_uids = [ + call.kwargs["packing_plan"].batch_uid + for call in c.student_policy.train.call_args_list + ] + assert observed_uids == [0, 1, 2, 3, 4, 5] + assert c.save_state["next_packing_batch_uid"] == 6 + + +def test_packing_batch_uid_accepts_checkpoint_without_high_water_mark(): + state = _default_off_policy_distillation_save_state() + del state["next_packing_batch_uid"] + state["total_steps"] = 7 + + batch_uids = _packing_batch_uids_from_state(state) + + assert [next(batch_uids), next(batch_uids)] == [14, 15] + assert state["next_packing_batch_uid"] == 16 + + +@pytest.mark.parametrize( + "fallback_metrics,error_match", + [ + ( + {"ipc_reconstruction_fallbacks": [1]}, + "missing consumer-side metrics", + ), + ( + { + "ipc_reconstruction_fallbacks": [2], + "ipc_reconstruction_fallbacks_t0": [1], + }, + "internally inconsistent", + ), + ], +) +def test_packed_dense_ipc_telemetry_fails_closed_and_releases_buffers( + mock_csd_components, fallback_metrics, error_match +): + c = mock_csd_components + _enable_lockstep_packing(c.master_config, global_batch_size=1, capacity=64) + c.master_config.distillation.update( + max_num_steps=1, + max_num_epochs=1, + val_period=0, + val_at_start=False, + val_at_end=False, + ) + c.student_policy.train.return_value = _make_train_results_with( + {"global_valid_toks": [10], "kl_loss": [0.3], **fallback_metrics} + ) + + with pytest.raises(RuntimeError, match=error_match): + _run_train(c) + + c.teacher_policy.release_ipc_buffer.assert_called_once_with() + + +# --------------------------------------------------------------------------- +# validate() loss-path branches +# --------------------------------------------------------------------------- + + +def _make_train_results_with(mb_metrics: dict) -> dict: + return { + "loss": torch.tensor(0.5), + "grad_norm": torch.tensor(1.0), + "all_mb_metrics": mb_metrics, + } + + +def test_validate_emits_pkl_metrics_only(mock_csd_components): + c = mock_csd_components + c.student_policy.train.return_value = _make_train_results_with( + {"kl_loss": [0.3], "ce_loss": [0.4]} + ) + + metrics, _timings = validate( + c.student_policy, + [c.teacher_policy], + c.val_dataloader, + c.loss_fn, + c.master_config, + skip_keys=xtoken_non_student_seq_keys(c.loss_fn), + ) + + assert "loss" in metrics + assert "kl_loss" in metrics + assert "ce_loss" in metrics + assert "kl_common" not in metrics + assert "l1_uncommon" not in metrics + + +def test_validate_collects_only_aggregate_metrics(mock_csd_components): + c = mock_csd_components + # validate summarizes only the aggregate loss / kl_loss / ce_loss / the + # aggregate gold path's kl_common / l1_uncommon. The per-teacher suffixed + # components (kl_common_t{i} / l1_uncommon_t{i}) are not collected here, so + # they don't appear in the val metrics. + c.student_policy.train.return_value = _make_train_results_with( + {"kl_common_t0": [0.2], "l1_uncommon_t0": [0.1], "ce_loss": [0.4]} + ) + + metrics, _timings = validate( + c.student_policy, + [c.teacher_policy], + c.val_dataloader, + c.loss_fn, + c.master_config, + skip_keys=xtoken_non_student_seq_keys(c.loss_fn), + ) + + assert "loss" in metrics + assert "ce_loss" in metrics + assert "kl_common" not in metrics + assert "l1_uncommon" not in metrics + assert "kl_loss" not in metrics + + +def test_validate_passes_ragged_packed_target_size_as_gbs(mock_csd_components): + c = mock_csd_components + _enable_lockstep_packing(c.master_config, global_batch_size=8, capacity=64) + c.student_policy.data_parallel_size = 2 + c.teacher_policy.data_parallel_size = 2 + ragged_batch = _make_batch(batch_size=3) + c.val_dataloader = MagicMock(spec=StatefulDataLoader) + c.val_dataloader.__iter__ = lambda self: iter([ragged_batch]) + + validate( + c.student_policy, + [c.teacher_policy], + c.val_dataloader, + c.loss_fn, + c.master_config, + skip_keys=xtoken_non_student_seq_keys(c.loss_fn), + ) + + train_kwargs = c.student_policy.train.call_args.kwargs + assert train_kwargs["gbs"] == 4 + assert len(train_kwargs["packing_plan"].canonical_batch_item_ids) == 4 + + +@pytest.mark.parametrize( + ("offload_student", "offload_optimizer", "restore_method", "unused_method"), + [ + (True, True, "offload_after_refit", "offload_before_refit"), + (False, True, "offload_before_refit", "offload_after_refit"), + ], +) +def test_validate_restores_student_between_teacher_state( + mock_csd_components, + offload_student, + offload_optimizer, + restore_method, + unused_method, +): + c = mock_csd_components + c.master_config.distillation["offload_student_after_step"] = offload_student + c.master_config.policy["offload_optimizer_for_logprob"] = offload_optimizer + + events = MagicMock() + events.attach_mock(c.student_policy, "student") + events.attach_mock(c.teacher_policy, "teacher") + validate( + c.student_policy, + [c.teacher_policy], + c.val_dataloader, + c.loss_fn, + c.master_config, + skip_keys=xtoken_non_student_seq_keys(c.loss_fn), + ) + + event_names = [call[0] for call in events.mock_calls] + train_indices = [ + index + for index, event_name in enumerate(event_names) + if event_name == "student.train" + ] + restore_indices = [ + index + for index, event_name in enumerate(event_names) + if event_name == f"student.{restore_method}" + ] + teacher_prepare_indices = [ + index + for index, event_name in enumerate(event_names) + if event_name == "teacher.prepare_for_lp_inference" + ] + + assert len(train_indices) == len(restore_indices) == 2 + assert train_indices[0] < restore_indices[0] < teacher_prepare_indices[1] + assert train_indices[1] < restore_indices[1] + assert getattr(c.student_policy, restore_method).call_count == 2 + getattr(c.student_policy, unused_method).assert_not_called() + + +# --------------------------------------------------------------------------- +# IPC buffer release-on-failure +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("num_teachers", [1, 2]) +def test_ipc_buffer_released_for_every_teacher_on_train_failure( + mock_csd_components, num_teachers +): + c = mock_csd_components + c.student_policy.train.side_effect = RuntimeError("boom") + c.master_config.distillation["max_num_steps"] = 1 + # No validation passes, so release is driven solely by the failing train + # step's finally (one release per teacher). + c.master_config.distillation["val_at_start"] = False + c.master_config.distillation["val_period"] = 0 + c.master_config.distillation["val_at_end"] = False + + # N same-vocab teachers: the release loop is teacher-type-agnostic, and + # same-vocab teachers reuse the student tokenization (the mock batch only + # carries teacher_0_* keys, so this stays valid for any N). + teacher_policies = [MagicMock(data_parallel_size=1) for _ in range(num_teachers)] + for t in teacher_policies: + t.get_full_logits_ipc.return_value = [{"payload_ipc": (4, 32)}] + c.loss_fn.num_teachers = num_teachers + c.loss_fn.projection_matrix_paths = [None] * num_teachers + # teacher_mbs is derived per entry in master_config.teachers, so it must + # have one entry per teacher policy. + c.master_config.teachers = [ + c.master_config.teachers[0] for _ in range(num_teachers) + ] + + with pytest.raises(RuntimeError): + xtoken_off_policy_distillation_train( + c.student_policy, + teacher_policies, + c.train_dataloader, + c.val_dataloader, + c.loss_fn, + c.logger, + c.checkpointer, + c.save_state, + c.master_config, + ) + + # The finally must release EVERY teacher's IPC buffer, not just the first + # (a misplaced finally / early break that frees only teacher 0 would pass a + # `>= 1` check while leaking the rest). + for t in teacher_policies: + assert t.release_ipc_buffer.call_count == 1 + + +# --------------------------------------------------------------------------- +# Multi-teacher wiring (CPU-only, no GPU loss math) +# --------------------------------------------------------------------------- + + +class _FakeLossFn(CSDLossFn): + """Minimal stand-in exposing the per-teacher metadata the trainer reads. + + Exposes the indexed metadata needed by the trainer and skip-key builder. + """ + + def __init__( + self, + projection_matrix_paths, + *, + teacher_topk_ipc_k=0, + teacher_vocab_sizes=None, + owner_local=False, + ): + self.num_teachers = len(projection_matrix_paths) + self.projection_matrix_paths = projection_matrix_paths + self.student_to_teacher_table_paths = [None] * self.num_teachers + self.teacher_vocab_sizes = teacher_vocab_sizes or [32] * self.num_teachers + self.cfg = CSDLossConfig( + **{ + "temperature": 1.0, + "kl_chunk_shift": False, + "csd_noise_filter_topk": 0, + "teacher_topk_ipc_k": teacher_topk_ipc_k, + "teacher_topk_ipc_support_mode": "row_topk", + "teacher_topk_ipc_keep_realized": True, + "owner_local": owner_local, + } + ) + + +def test_skip_keys_builder_cross_and_same_vocab(): + # teacher 0 cross-tokenizer (full logits); teacher 1 same-vocab. + loss_fn = _FakeLossFn(projection_matrix_paths=["/p0.pt", None]) + keys = xtoken_non_student_seq_keys(loss_fn) + # Cross-tokenizer teacher 0: IPC handle list + teacher tokens + the + # teacher-seq / max_pairs alignment keys are skipped. + assert "teacher_0_full_logits_ipc" in keys + assert "teacher_0_sparse_logits_ipc" in keys + assert "teacher_0_input_ids" in keys + assert "teacher_0_token_mask" in keys + assert "alignment_0_pair_valid" in keys + assert "alignment_0_teacher_chunk_id" in keys + # Student-seq alignment keys ([B, T_s]) and num_chunks ([B]) are NOT skipped. + assert "alignment_0_student_chunk_id" not in keys + assert "alignment_0_num_chunks" not in keys + # Same-vocab teacher 1 also ships full logits over IPC, so its handle-list + # key (a non-tensor) is skipped; it reuses the student tokenization, so it + # has no teacher-seq token keys and no teacher-indexed alignment keys. + assert "teacher_1_full_logits_ipc" in keys + assert "teacher_1_sparse_logits_ipc" in keys + assert not any(k.startswith("alignment_1_") for k in keys) + assert "teacher_1_input_ids" not in keys + + +def test_reduce_mb_metric_means_per_chunk_diagnostics_and_sums_shares(): + """CSD per-chunk diagnostics are per-mb averages: mean-reduce, never sum. + + Summing them scales the reported value with the microbatch count — the + defect that produced a ``top1_acc_per_chunk`` of 110.60 on a 128-microbatch + step. The KD/CE terms are normalized by a *global* denominator in the loss, + so those stay sum-reduced, as do the ``num_*`` counts. + """ + n_mb = 128 + acc = [0.8] * n_mb + # Mean-reduced, both bare and with the per-teacher ``_t{i}`` suffix. + for key in ( + "top1_acc_per_chunk", + "kl_common_per_chunk", + "kl_partition_last_per_chunk", + ): + assert reduce_mb_metric(key, acc) == pytest.approx(0.8) + assert reduce_mb_metric(f"{key}_t0", acc) == pytest.approx(0.8) + assert reduce_mb_metric(f"{key}_t11", acc) == pytest.approx(0.8) + assert reduce_mb_metric("kl_loss_scale", [4.0] * n_mb) == pytest.approx(4.0) + # Sum-reduced: global-denominator loss shares and counts, suffixed or not. + assert reduce_mb_metric("kl_loss", [0.5] * 4) == pytest.approx(2.0) + assert reduce_mb_metric("kl_loss_t0", [0.5] * 4) == pytest.approx(2.0) + assert reduce_mb_metric("ce_loss", [0.5] * 4) == pytest.approx(2.0) + assert reduce_mb_metric("num_common_chunks_t0", [10] * 4) == pytest.approx(40) + assert reduce_mb_metric("teacher_0/routed_samples", [2, 3]) == pytest.approx(5) + assert reduce_mb_metric("teacher_0/routed_tokens", [7, 11]) == pytest.approx(18) + assert reduce_mb_metric("teacher_0/weighted_kl", [0.2, 0.3]) == pytest.approx(0.5) + + +def test_dense_ipc_telemetry_counts_tp2_cp2_fallbacks_once_per_logical_row(capsys): + """TP/CP-replicated worker results must not multiply the logical count. + + The packed TP2/CP2 policy returns one replicated result per DP rank, and + each of the four DP ranks contributes 32 logical MBS1 records. A TP2 + teacher cannot use the full-vocabulary zero-copy path, so all 128 records + report one actual reconstruction fallback for teacher 0. + """ + per_dp_records = [[1] * 32 for _ in range(4)] + model_parallel_deduplicated = [ + value for dp_records in per_dp_records for value in dp_records + ] + + csd_mod._log_dense_ipc_reconstruction_telemetry( + { + "ipc_reconstruction_fallbacks": model_parallel_deduplicated, + "ipc_reconstruction_fallbacks_t0": model_parallel_deduplicated, + }, + packing_plan=SimpleNamespace(batch_uid=17), + num_teachers=1, + ) + + assert capsys.readouterr().out == ( + "CSD_IPC_RECONSTRUCTION batch_uid=17 " + "reconstruction_fallbacks=128 teacher_0=128\n" + ) + + +def test_skip_keys_builder_same_vocab_full_logits(): + loss_fn = _FakeLossFn(projection_matrix_paths=[None]) + # Same-vocab full-logits teacher: only the IPC handle list is skipped. + assert xtoken_non_student_seq_keys(loss_fn) == frozenset( + {"teacher_0_full_logits_ipc", "teacher_0_sparse_logits_ipc"} + ) + + +def test_export_teacher_logits_packs_indexed_keys_and_runs_serially(): + # teacher 0 cross-tokenizer; teacher 1 same-vocab. Both ship full-vocab + # logits over the unified get_full_logits_ipc producer (always-full). + loss_fn = _FakeLossFn(projection_matrix_paths=["/p0.pt", None]) + # Attach both teacher mocks to one parent so their calls land in a single + # ordered list (lets us assert the serial interleaving below). Configure + # the child return values after attaching. + parent = MagicMock() + t0 = MagicMock() + t1 = MagicMock() + parent.attach_mock(t0, "t0") + parent.attach_mock(t1, "t1") + t0.get_full_logits_ipc.return_value = [{"payload_ipc": 0}] + t1.get_full_logits_ipc.return_value = [{"payload_ipc": 1}] + # Cross-tokenizer teacher 0 needs teacher_0_*/alignment_0_*; same-vocab + # teacher 1 reuses the student tokenization (no teacher_1_* token keys). + batch = _make_batch(num_teachers=1) + + train_data = export_teacher_logits_and_pack( + [t0, t1], loss_fn, batch, teacher_mbs=[1, 1] + ) + + # Cross-tokenizer teacher 0 -> full-logits IPC + its alignment payload. + assert "teacher_0_full_logits_ipc" in train_data + assert "alignment_0_pair_valid" in train_data + assert "teacher_0_input_ids" in train_data + # Same-vocab teacher 1 -> full-logits IPC handle list, no dense tensors, + # no projection/alignment keys (it reuses the student tokenization). + assert "teacher_1_full_logits_ipc" in train_data + assert "teacher_1_input_ids" not in train_data + assert "alignment_1_pair_valid" not in train_data + # Both teachers ship full logits via get_full_logits_ipc with their own MBS. + t0.get_full_logits_ipc.assert_called_once() + assert t0.get_full_logits_ipc.call_args.kwargs["micro_batch_size"] == 1 + t1.get_full_logits_ipc.assert_called_once() + assert t1.get_full_logits_ipc.call_args.kwargs["micro_batch_size"] == 1 + # Serial collocated execution: each teacher onloaded then offloaded, AND + # teacher 0 is offloaded before teacher 1 is onloaded (never both resident). + for t in (t0, t1): + t.prepare_for_lp_inference.assert_called_once() + t.offload_after_refit.assert_called_once() + call_names = [c[0] for c in parent.mock_calls] + assert call_names.index("t0.offload_after_refit") < call_names.index( + "t1.prepare_for_lp_inference" + ) + + +def test_export_teacher_logits_preserves_student_routed_experts(): + teacher = MagicMock() + teacher.get_full_logits_ipc.return_value = [{"payload_ipc": 0}] + loss_fn = _FakeLossFn(projection_matrix_paths=["/projection.pt"]) + batch = _make_batch(num_teachers=1) + routed_experts = torch.arange( + batch["input_ids"].shape[0] * batch["input_ids"].shape[1] * 2 * 2 + ).reshape(batch["input_ids"].shape[0], batch["input_ids"].shape[1], 2, 2) + batch["routed_experts"] = routed_experts + + train_data = export_teacher_logits_and_pack( + [teacher], loss_fn, batch, teacher_mbs=[1] + ) + + assert train_data["routed_experts"] is routed_experts + + +def test_export_teacher_logits_threads_lockstep_plan_and_occurrence_ids(): + cfg = _make_master_config() + _enable_lockstep_packing(cfg, global_batch_size=2, capacity=64) + batch = _make_batch(batch_size=2, t_student=4, t_teacher=6) + batch["student_semantic_regions"] = [ + ((0, "assistant", "content", 1, 3),), + ((2, "assistant", "eot", 3, 4),), + ] + batch["teacher_0_semantic_regions"] = [ + ((0, "assistant", "content", 2, 4),), + ((2, "assistant", "eot", 4, 5),), + ] + plan = build_csd_lockstep_packing_plan( + batch, cfg, batch_uid=17, data_parallel_size=1 + ) + assert plan is not None + teacher = MagicMock() + teacher.get_full_logits_ipc.return_value = [ + {"batch_item_id": item_id, "teacher_shards": []} + for item_id in plan.canonical_batch_item_ids + ] + loss_fn = _FakeLossFn(projection_matrix_paths=["/projection.pt"]) + + train_data = export_teacher_logits_and_pack( + [teacher], + loss_fn, + batch, + teacher_mbs=[1], + packing_plan=plan, + ) + + assert train_data["sample_id"] == batch["sample_id"] + assert torch.equal(train_data["batch_item_id"], batch["batch_item_id"]) + first_occurrence = plan.canonical_batch_item_ids[0] + assert train_data["student_semantic_regions"][0][0] == ( + first_occurrence, + 0, + "assistant", + "content", + 1, + 3, + ) + assert train_data["teacher_0_semantic_regions"][0][0][0] == first_occurrence + teacher_data = teacher.get_full_logits_ipc.call_args.args[0] + assert torch.equal(teacher_data["batch_item_id"], batch["batch_item_id"]) + assert teacher.get_full_logits_ipc.call_args.kwargs["packing_plan"] is plan + assert ( + teacher.get_full_logits_ipc.call_args.kwargs["packing_side_id"] == "teacher_0" + ) + + +def test_build_teacher_force_ids_respects_position_zero_and_shift(): + batch = _make_batch(batch_size=1, t_student=5, t_teacher=5) + batch["teacher_0_input_ids"] = torch.tensor([[10, 11, 12, 13, 14]]) + batch["alignment_0_pair_valid"] = torch.tensor([[True, True, False]]) + batch["alignment_0_student_chunk_id"] = torch.tensor([[0, 0, -1, 1, 1]]) + batch["alignment_0_teacher_chunk_id"] = torch.tensor([[0, 0, 1, 1, -1]]) + loss_config = { + "kl_chunk_shift": True, + "csd_noise_filter_topk": 0, + "csd_pure_alm": False, + "csd_realized_only_mismatch": False, + "teacher_topk_ipc_k": 8192, + "teacher_topk_ipc_keep_realized": True, + } + + force_ids = _build_teacher_force_include_token_ids( + batch, + teacher_idx=0, + loss_config=loss_config, + ) + + # Chunk 0 starts at position zero, so it remains unshifted. Chunk 1 starts + # later and its labels at teacher positions 2,3 are attached to predictors + # 1,2, with the latter write taking precedence at position 1. + assert torch.equal(force_ids, torch.tensor([[10, 12, 13, -1, -1]])) + + +@pytest.mark.parametrize( + "kl_chunk_shift,expected", + [ + # Chunk 0 starts at position zero and stays unshifted; chunk 1's labels + # at teacher positions 2,3 attach to predictors 1,2, so row 1 keeps + # both its unshifted label 11 and its shifted label 12. + (True, [[10, -1], [11, 12], [13, -1], [-1, -1], [-1, -1]]), + (False, [[10, -1], [11, -1], [12, -1], [13, -1], [-1, -1]]), + ], +) +def test_build_teacher_force_id_slots_keep_both_labels_of_a_shared_row( + kl_chunk_shift, expected +): + batch = _make_batch(batch_size=1, t_student=5, t_teacher=5) + batch["teacher_0_input_ids"] = torch.tensor([[10, 11, 12, 13, 14]]) + batch["alignment_0_pair_valid"] = torch.tensor([[True, True, False]]) + batch["alignment_0_student_chunk_id"] = torch.tensor([[0, 0, -1, 1, 1]]) + batch["alignment_0_teacher_chunk_id"] = torch.tensor([[0, 0, 1, 1, -1]]) + + force_ids = _build_teacher_force_include_token_id_slots( + batch, teacher_idx=0, kl_chunk_shift=kl_chunk_shift + ) + + assert force_ids.shape == (1, 5, csd_mod.SPARSE_TEACHER_FORCE_WIDTH) + assert force_ids.tolist() == [expected] + + +def test_build_teacher_force_id_slots_skip_masked_and_one_sided_chunks(): + batch = _make_batch(batch_size=2, t_student=4, t_teacher=4) + batch["teacher_0_input_ids"] = torch.tensor([[10, 11, 12, 13], [20, 21, 22, 23]]) + batch["sample_mask"] = torch.tensor([1, 0]) + batch["alignment_0_pair_valid"] = torch.tensor([[True, True], [True, True]]) + # Chunk 1 of sample 0 has no student positions, so its span is [0, 0). + batch["alignment_0_student_chunk_id"] = torch.tensor([[0, 0, -1, -1], [0, 1, 1, 1]]) + batch["alignment_0_teacher_chunk_id"] = torch.tensor([[0, 1, 1, -1], [0, 1, 1, 1]]) + + force_ids = _build_teacher_force_include_token_id_slots( + batch, teacher_idx=0, kl_chunk_shift=True + ) + + assert force_ids[0].tolist() == [[10, -1], [-1, -1], [-1, -1], [-1, -1]] + assert bool((force_ids[1] == -1).all()) + + +@pytest.mark.parametrize("realized_only", [False, True]) +def test_build_teacher_force_ids_follow_realized_only_mismatch(realized_only): + """Realized-only mismatch keeps realized labels even without keep_realized.""" + batch = _make_batch(batch_size=1, t_student=5, t_teacher=5) + batch["teacher_0_input_ids"] = torch.tensor([[10, 11, 12, 13, 14]]) + batch["alignment_0_pair_valid"] = torch.tensor([[True, True, False]]) + batch["alignment_0_student_chunk_id"] = torch.tensor([[0, 0, -1, 1, 1]]) + batch["alignment_0_teacher_chunk_id"] = torch.tensor([[0, 0, 1, 1, -1]]) + loss_config = { + "kl_chunk_shift": True, + "csd_noise_filter_topk": 0, + "csd_pure_alm": False, + "csd_realized_only_mismatch": realized_only, + "teacher_topk_ipc_k": 8192, + "teacher_topk_ipc_keep_realized": False, + } + + force_ids = _build_teacher_force_include_token_ids( + batch, + teacher_idx=0, + loss_config=loss_config, + ) + + if realized_only: + assert torch.equal(force_ids, torch.tensor([[10, 12, 13, -1, -1]])) + else: + assert force_ids is None + + +def test_owner_local_export_routes_sparse_teachers_to_the_megatron_sender(): + loss_fn = _FakeLossFn( + projection_matrix_paths=["/p0.pt", None], + teacher_topk_ipc_k=4, + teacher_vocab_sizes=[151669, 128256], + owner_local=True, + ) + t0 = MagicMock() + t1 = MagicMock() + t0.get_sparse_logits_ipc.return_value = [{"teacher_shards": ["sparse"]}] + t1.get_full_logits_ipc.return_value = [{"teacher_shards": ["dense"]}] + batch = _make_batch(batch_size=2, num_teachers=1) + batch["teacher_0_input_ids"] = torch.tensor([[3, 4, 5, 6], [7, 8, 9, 10]]) + + train_data = export_teacher_logits_and_pack( + [t0, t1], loss_fn, batch, teacher_mbs=[1, 2] + ) + + assert train_data["teacher_0_sparse_logits_ipc"] == [{"teacher_shards": ["sparse"]}] + assert train_data["teacher_1_full_logits_ipc"] == [{"teacher_shards": ["dense"]}] + t0.get_topk_logits_ipc.assert_not_called() + call = t0.get_sparse_logits_ipc.call_args + assert call.kwargs == { + "k": 4, + "temperature": 1.0, + "real_vocab_size": 151669, + "noise_filter_k": 0, + "micro_batch_size": 1, + "timer": None, + } + force = call.args[0]["force_include_token_ids"] + expected = _build_teacher_force_include_token_id_slots( + batch, teacher_idx=0, kl_chunk_shift=False + ) + assert force.shape == (2, 4, csd_mod.SPARSE_TEACHER_FORCE_WIDTH) + assert torch.equal(force, expected) + assert force[0, 0].tolist() == [3, -1] + + +def test_owner_local_dense_export_keeps_full_logits(): + loss_fn = _FakeLossFn( + projection_matrix_paths=["/p0.pt"], teacher_topk_ipc_k=0, owner_local=True + ) + teacher = MagicMock() + teacher.get_full_logits_ipc.return_value = [{"teacher_shards": ["dense"]}] + + train_data = export_teacher_logits_and_pack( + [teacher], loss_fn, _make_batch(), teacher_mbs=[1] + ) + + assert train_data["teacher_0_full_logits_ipc"] == [{"teacher_shards": ["dense"]}] + teacher.get_sparse_logits_ipc.assert_not_called() + teacher.get_topk_logits_ipc.assert_not_called() + + +def _owner_local_master_config(*, k: int = 0, teacher_backend: str = "megatron"): + cfg = _make_master_config() + cfg.policy["automodel_cfg"]["enabled"] = False + cfg.policy["megatron_cfg"] = _megatron_cfg(tp=2, cp=2) + if teacher_backend == "megatron": + cfg.teachers[0].automodel_cfg["enabled"] = False + cfg.teachers[0].megatron_cfg = { + **_megatron_cfg(tp=2, cp=2), + "sequence_parallel": False, + } + cfg.loss_fn.teacher_topk_ipc_k = k + return cfg + + +@pytest.mark.parametrize( + "change,expected", + [ + (None, []), + ("automodel_student", ["the student policy is not Megatron"]), + ("pp2", ["policy pipeline_model_parallel_size=2 (needs 1)"]), + ("packing", ["policy sequence packing is enabled"]), + ("dynamic_batching", ["policy dynamic batching is enabled"]), + ("pure_alm", ["csd_pure_alm=true"]), + ("averaged_logits", ["kd_loss_mode='averaged_logits' (needs 'sum')"]), + ("sparse", []), + ("sparse_automodel_teacher", ["teachers[0] is not Megatron"]), + ("sparse_odd_divisor", ["is not a multiple of 4"]), + ("sparse_sequence_parallel", ["is not a multiple of 12"]), + ], +) +def test_owner_local_eligibility_truth_table(change, expected): + cfg = _owner_local_master_config() + if change == "automodel_student": + cfg.policy["automodel_cfg"]["enabled"] = True + cfg.policy["megatron_cfg"]["enabled"] = False + elif change == "pp2": + cfg.policy["megatron_cfg"]["pipeline_model_parallel_size"] = 2 + elif change == "packing": + cfg.policy["sequence_packing"] = {"enabled": True} + elif change == "dynamic_batching": + cfg.policy["dynamic_batching"] = {"enabled": True} + elif change == "pure_alm": + cfg.loss_fn.csd_pure_alm = True + elif change == "averaged_logits": + cfg.loss_fn.kd_loss_mode = "averaged_logits" + elif change is not None and change.startswith("sparse"): + cfg.loss_fn.teacher_topk_ipc_k = 128 + if change == "sparse_automodel_teacher": + cfg.teachers[0].automodel_cfg["enabled"] = True + cfg.teachers[0].megatron_cfg = _megatron_cfg(enabled=False) + elif change == "sparse_odd_divisor": + cfg.teachers[0].make_sequence_length_divisible_by = 6 + elif change == "sparse_sequence_parallel": + cfg.teachers[0].megatron_cfg = { + **_megatron_cfg(tp=3, cp=2), + "sequence_parallel": True, + } + + reasons = csd_mod.csd_owner_local_ineligibility(cfg, cfg.loss_fn) + + assert len(reasons) == len(expected) + for reason, fragment in zip(reasons, expected): + assert fragment in reason + + +def test_owner_local_ignores_same_tokenizer_teachers_for_sparse_checks(): + cfg = _owner_local_master_config(k=128, teacher_backend="automodel") + cfg.teachers[0].aligner.student_to_teacher_table_path = None + cfg.teachers[0].aligner.teacher_to_student_table_path = None + assert csd_mod.csd_owner_local_ineligibility(cfg, cfg.loss_fn) == [] + + +@pytest.mark.parametrize( + "k,teacher_backend,pure_alm,expected_owner_local", + [ + (0, "megatron", False, True), + (0, "automodel", False, True), + (128, "megatron", False, True), + (0, "megatron", True, False), + ], +) +def test_setup_injects_owner_local_for_eligible_megatron_students( + k, teacher_backend, pure_alm, expected_owner_local +): + cfg = _owner_local_master_config(k=k, teacher_backend=teacher_backend) + cfg.loss_fn.csd_pure_alm = pure_alm + + _, mocks = _patched_setup_call(cfg) + + assert mocks["loss"].call_args.args[0]["owner_local"] is expected_owner_local + + +def test_setup_keeps_automodel_students_replicated(): + _, mocks = _patched_setup_call(_make_master_config()) + assert mocks["loss"].call_args.args[0]["owner_local"] is False + + +def test_setup_rejects_sparse_ipc_a_megatron_student_cannot_own(): + cfg = _owner_local_master_config(k=128, teacher_backend="automodel") + with pytest.raises( + ValueError, + match=r"needs owner-local CSD evaluation.*teachers\[0\] is not Megatron" + r".*teacher_topk_ipc_k=0", + ): + _patched_setup_call(cfg) + cfg = _owner_local_master_config(k=128) + cfg.loss_fn.csd_pure_alm = True + with pytest.raises(ValueError, match="csd_pure_alm=true"): + _patched_setup_call(cfg) + + +def test_setup_rejects_megatron_sparse_teachers_for_automodel_students(): + cfg = _make_master_config() + cfg.loss_fn.teacher_topk_ipc_k = 128 + cfg.teachers[0].automodel_cfg["enabled"] = False + cfg.teachers[0].megatron_cfg = _megatron_cfg() + with pytest.raises( + ValueError, match="Automodel student requires Automodel cross-tokenizer" + ): + _patched_setup_call(cfg) + + +def test_export_teacher_logits_mixes_sparse_cross_and_dense_same_vocab(): + loss_fn = _FakeLossFn( + projection_matrix_paths=["/p0.pt", None], + teacher_topk_ipc_k=8192, + teacher_vocab_sizes=[151669, 128256], + ) + t0 = MagicMock() + t1 = MagicMock() + t0.get_topk_logits_ipc.return_value = [{"teacher_shards": ["sparse"]}] + t1.get_full_logits_ipc.return_value = [{"teacher_shards": ["dense"]}] + batch = _make_batch(num_teachers=1) + + train_data = export_teacher_logits_and_pack( + [t0, t1], loss_fn, batch, teacher_mbs=[1, 2] + ) + + assert train_data["teacher_0_sparse_logits_ipc"] == [{"teacher_shards": ["sparse"]}] + assert train_data["teacher_1_full_logits_ipc"] == [{"teacher_shards": ["dense"]}] + assert "teacher_0_full_logits_ipc" not in train_data + assert "teacher_1_sparse_logits_ipc" not in train_data + t0.get_full_logits_ipc.assert_not_called() + t1.get_topk_logits_ipc.assert_not_called() + sparse_kwargs = t0.get_topk_logits_ipc.call_args.kwargs + assert sparse_kwargs["k"] == 8192 + assert sparse_kwargs["temperature"] == 1.0 + assert sparse_kwargs["vocab_size"] == 151669 + assert sparse_kwargs["micro_batch_size"] == 1 + assert sparse_kwargs["support_mode"] == "row_topk" + assert sparse_kwargs["gt_filter_topk"] is None + assert torch.equal( + t0.get_topk_logits_ipc.call_args.args[0]["force_include_token_ids"], + torch.zeros((1, 4), dtype=torch.long), + ) + assert t1.get_full_logits_ipc.call_args.kwargs["micro_batch_size"] == 2 + + +def test_setup_preserves_aligner_config_across_interleaved_teacher_types(): + cfg = _make_master_config() + # Interleave a same-vocab teacher between two cross-tokenizer teachers. + cfg.teachers.append( + TeacherConfig( + **{ + "aligner": {"projection_matrix_path": None}, + "weight": 0.5, + "automodel_cfg": { + "enabled": True, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + }, + "max_total_sequence_length": 64, + "make_sequence_length_divisible_by": 8, + "train_global_batch_size": 1, + "train_micro_batch_size": 1, + "tokenizer": {"name": "student-tok"}, + } + ) + ) + cfg.teachers.append( + TeacherConfig( + **{ + "aligner": { + "projection_matrix_path": None, + "student_to_teacher_table_path": "/tmp/forward-2.pt", + "teacher_to_student_table_path": "/tmp/reverse-2.pt", + "drop_first_assistant_chunk_kl": True, + }, + "weight": 0.25, + "automodel_cfg": { + "enabled": True, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + }, + "max_total_sequence_length": 64, + "make_sequence_length_divisible_by": 8, + "train_global_batch_size": 1, + "train_micro_batch_size": 1, + "tokenizer": {"name": "teacher-2-tok"}, + } + ) + ) + student_tok = _make_tokenizer(32) + teacher_toks = [_make_tokenizer(24), _make_tokenizer(32), _make_tokenizer(28)] + train_ds = MagicMock() + train_ds.__len__ = MagicMock(return_value=4) + + with ( + patch.object(csd_mod, "RayVirtualCluster") as mock_cluster, + patch.object(csd_mod, "Policy") as mock_policy_cls, + patch.object(csd_mod, "Logger"), + patch.object(csd_mod, "CheckpointManager") as mock_cp_cls, + patch.object(csd_mod, "TokenAligner") as mock_aligner_cls, + patch.object(csd_mod, "CSDCrossTokenizerCollator") as mock_collator_cls, + patch.object(csd_mod, "CSDLossFn") as mock_loss_cls, + patch.object(csd_mod, "StatefulDataLoader") as mock_dl_cls, + patch.object(csd_mod, "assert_teacher_student_batch_grid"), + patch.object(csd_mod, "assert_xtoken_ipc_node_local"), + ): + mock_cp_cls.return_value.get_latest_checkpoint_path.return_value = None + mock_cp_cls.return_value.load_training_info.return_value = None + mock_cp_cls.return_value.get_resume_paths.return_value = (None, None) + mock_dl_cls.side_effect = lambda *a, **kw: MagicMock(spec=StatefulDataLoader) + mock_policy_cls.side_effect = lambda *a, **kw: MagicMock(data_parallel_size=1) + + (_student, teachers, *_rest) = setup( + cfg, + student_tokenizer=student_tok, + teacher_tokenizers=teacher_toks, + train_dataset=train_ds, + val_dataset=None, + ) + + # One teacher Policy per entry (+ the student), and the colocation cap + # accounts for all teacher groups + the student. + assert isinstance(teachers, list) and len(teachers) == 3 + assert mock_policy_cls.call_count == 4 # 3 teachers + 1 student + assert mock_cluster.call_args.kwargs["max_colocated_worker_groups"] == 4 + assert mock_aligner_cls.call_count == 2 + aligners = mock_collator_cls.call_args.kwargs["aligners"] + assert aligners[0] is not None and aligners[1] is None and aligners[2] is not None + assert mock_collator_cls.call_args.kwargs[ + "drop_first_assistant_chunk_kl_by_teacher" + ] == [False, False, True] + # Per-teacher metadata injected as parallel lists. + injected_cfg = mock_loss_cls.call_args.args[0] + assert injected_cfg["projection_matrix_paths"] == [None, None, None] + assert injected_cfg["student_to_teacher_table_paths"] == [ + "/tmp/forward.pt", + None, + "/tmp/forward-2.pt", + ] + assert injected_cfg["teacher_to_student_table_paths"] == [ + "/tmp/reverse.pt", + None, + "/tmp/reverse-2.pt", + ] + assert injected_cfg["teacher_weights"] == [1.0, 0.5, 0.25] + assert injected_cfg["teacher_vocab_sizes"] == [24, 32, 28] + + +def test_setup_rejects_same_vocab_teacher_with_mismatched_vocab(): + # A null-projection (same-vocab) teacher whose vocab != the student's is a + # config error — same-vocab detection is by null projection path, but the + # teacher must actually share the student vocab for the direct KL. Caught + # in setup() (the right place — it has the real tokenizers; tokenizer + # *names* would wrongly flag Llama-3.2-3B vs -1B, which share a vocab). + cfg = _make_master_config() + cfg.teachers[0].aligner.projection_matrix_path = None + cfg.teachers[0].aligner.student_to_teacher_table_path = None + cfg.teachers[0].aligner.teacher_to_student_table_path = None # mark same-vocab + student_tok = _make_tokenizer(32) + teacher_toks = [_make_tokenizer(24)] # 24 != 32 -> mismatch + with ( + patch.object(csd_mod, "RayVirtualCluster") as mock_cluster, + pytest.raises(ValueError, match="safe student-token reuse"), + ): + setup( + cfg, + student_tokenizer=student_tok, + teacher_tokenizers=teacher_toks, + train_dataset=MagicMock(), + val_dataset=None, + ) + # The check fires before any cluster/policy construction. + assert mock_cluster.call_count == 0 + + +def test_same_tokenizer_reuse_rejects_equal_size_different_mapping(): + cfg = _make_master_config() + cfg.teachers[0].aligner.projection_matrix_path = None + cfg.teachers[0].aligner.student_to_teacher_table_path = None + cfg.teachers[0].aligner.teacher_to_student_table_path = None + student = _make_tokenizer(4) + teacher = _make_tokenizer(4) + teacher.get_vocab.return_value = { + "token-1": 0, + "token-0": 1, + "token-2": 2, + "token-3": 3, + } + + with pytest.raises(ValueError, match="differing fields: .*vocab"): + validate_csd_tokenizer_reuse(cfg, student, [teacher]) + + +def test_same_tokenizer_reuse_rejects_template_kwargs_difference(): + cfg = _make_master_config() + cfg.teachers[0].aligner.projection_matrix_path = None + cfg.teachers[0].aligner.student_to_teacher_table_path = None + cfg.teachers[0].aligner.teacher_to_student_table_path = None + cfg.policy["tokenizer"]["chat_template_kwargs"] = {"enable_thinking": False} + cfg.teachers[0].tokenizer["chat_template_kwargs"] = {"enable_thinking": True} + + with pytest.raises(ValueError, match="chat_template_kwargs"): + validate_csd_tokenizer_reuse(cfg, _make_tokenizer(4), [_make_tokenizer(4)]) + + +def test_same_tokenizer_reuse_rejects_slow_tokenizer_even_with_same_vocab(): + cfg = _make_master_config() + cfg.teachers[0].aligner.projection_matrix_path = None + cfg.teachers[0].aligner.student_to_teacher_table_path = None + cfg.teachers[0].aligner.teacher_to_student_table_path = None + student = _make_tokenizer(4) + teacher = _make_tokenizer(4) + teacher.is_fast = False + + with pytest.raises(ValueError, match="slow-tokenizer.*cannot be proven"): + validate_csd_tokenizer_reuse(cfg, student, [teacher]) + + +def test_same_tokenizer_reuse_rejects_different_fast_backend_behavior(): + cfg = _make_master_config() + cfg.teachers[0].aligner.projection_matrix_path = None + cfg.teachers[0].aligner.student_to_teacher_table_path = None + cfg.teachers[0].aligner.teacher_to_student_table_path = None + student = _make_tokenizer(4) + teacher = _make_tokenizer(4) + teacher.backend_tokenizer.to_str.return_value = "different-normalizer" + + with pytest.raises(ValueError, match="differing fields: .*backend"): + validate_csd_tokenizer_reuse(cfg, student, [teacher]) + + +def test_lockstep_setup_accepts_supported_automodel_cp1(): + cfg = _make_master_config() + _enable_lockstep_packing(cfg) + + assert validate_csd_packing_setup(cfg) == 1 + + +def test_lockstep_setup_rejects_renamed_mamba_from_architecture_metadata( + monkeypatch, +): + cfg = _make_master_config() + _enable_lockstep_packing(cfg) + cfg.policy["model_name"] = "/models/innocuous-local-name" + mamba_config = SimpleNamespace( + model_type="nemotron_h", + architectures=["NemotronHForCausalLM"], + to_dict=lambda: { + "model_type": "nemotron_h", + "architectures": ["NemotronHForCausalLM"], + "hybrid_override_pattern": "M*-M*-", + "mamba_d_state": 128, + }, + ) + monkeypatch.setattr( + csd_mod, "_load_csd_packed_model_config", lambda config: mamba_config + ) + + with pytest.raises(ValueError, match="Nano/Mamba recurrent state isolation"): + validate_csd_packing_setup(cfg) + + +@pytest.mark.parametrize( + ("mutate", "message"), + [ + ( + lambda cfg: cfg.teachers[0].sequence_packing.update(enabled=False), + "student and every teacher", + ), + ( + lambda cfg: cfg.policy["dynamic_batching"].update(enabled=True), + "dynamic batching", + ), + ( + lambda cfg: cfg.data["train"].update(characters_per_sample=128), + "characters_per_sample=null", + ), + ( + lambda cfg: cfg.data.update(validation={"characters_per_sample": 128}), + "data.validation.characters_per_sample=null", + ), + ( + lambda cfg: setattr(cfg.collator, "num_packed_rows", 2), + "num_packed_rows must be 1", + ), + ( + lambda cfg: cfg.policy["sequence_packing"].update( + algorithm="modified_first_fit_decreasing" + ), + "lockstep_first_fit_decreasing", + ), + ( + lambda cfg: cfg.policy["sequence_packing"].pop("fuse_loss"), + "fuse_loss must be set explicitly", + ), + ( + lambda cfg: setattr(cfg.loss_fn, "teacher_topk_ipc_k", 8), + "dense teacher IPC only", + ), + ( + lambda cfg: setattr(cfg.loss_fn, "kd_loss_mode", "select_teacher"), + "static additive teacher aggregation", + ), + ( + lambda cfg: cfg.policy["automodel_cfg"].update(context_parallel_size=2), + "Automodel CP=1 only", + ), + ( + lambda cfg: cfg.policy["automodel_cfg"].update(sequence_parallel=True), + "sequence parallelism", + ), + ], +) +def test_lockstep_setup_rejects_unsupported_modes(mutate, message): + cfg = _make_master_config() + _enable_lockstep_packing(cfg) + mutate(cfg) + + with pytest.raises(ValueError, match=message): + validate_csd_packing_setup(cfg) + + +def test_lockstep_setup_resolves_character_grouping_for_validation_defaults(): + cfg = _make_master_config() + _enable_lockstep_packing(cfg) + cfg.data["train"]["characters_per_sample"] = None + cfg.data["default"] = {"characters_per_sample": 128} + cfg.data["validation"] = {"dataset_name": "validation"} + + with pytest.raises(ValueError, match="data.validation.characters_per_sample=null"): + validate_csd_packing_setup(cfg) + + +def test_lockstep_plan_assigns_unique_occurrence_ids_after_validation_padding(): + cfg = _make_master_config() + cfg.cluster = {"num_nodes": 1, "gpus_per_node": 2} + cfg.policy["automodel_cfg"]["tensor_parallel_size"] = 1 + cfg.teachers[0].automodel_cfg["tensor_parallel_size"] = 1 + _enable_lockstep_packing(cfg, global_batch_size=4, capacity=64) + batch = csd_mod.pad_distillation_val_batch(_make_batch(batch_size=3), 4) + + plan = build_csd_lockstep_packing_plan( + batch, + cfg, + batch_uid=7, + data_parallel_size=2, + ) + + assert plan is not None + assert batch["sample_id"] == ["sample-0", "sample-1", "sample-2", "sample-2"] + assert len(set(batch["batch_item_id"].tolist())) == 4 + assert plan.canonical_batch_item_ids == tuple(batch["batch_item_id"].tolist()) + assert plan.sides["student"].raw_lengths == (4, 4, 4, 4) + assert plan.sides["teacher_0"].raw_lengths == (4, 4, 4, 4) + assert plan.sides["student"].rank_bin_indices == ((0,), (1,)) + assert plan.bins == ( + tuple(plan.canonical_batch_item_ids[:2]), + tuple(plan.canonical_batch_item_ids[2:]), + ) + + +def test_dtensor_tp_student_plan_preserves_fixed_training_tail(): + cfg = _make_master_config() + _enable_lockstep_packing(cfg, global_batch_size=2, capacity=64) + cfg.policy["automodel_cfg"]["tensor_parallel_size"] = 2 + batch = _make_batch(batch_size=2, t_student=4, t_teacher=6) + + plan = build_csd_lockstep_packing_plan( + batch, + cfg, + batch_uid=8, + data_parallel_size=1, + ) + + assert plan is not None + assert len(plan.bins) == 1 + assert plan.sides["student"].physical_tokens_by_bin == (64,) + assert plan.sides["student"].padded_cu_seqlens_by_bin[0][-1] == 64 + # The inference-only teacher does not inherit the student's TP training + # tail; its side-local physical geometry remains independently minimal. + assert plan.sides["teacher_0"].physical_tokens_by_bin == (16,) + + +def test_lockstep_telemetry_proves_multi_sample_bins(capsys): + cfg = _make_master_config() + _enable_lockstep_packing(cfg, global_batch_size=2, capacity=64) + batch = _make_batch(batch_size=2, t_student=4, t_teacher=6) + plan = build_csd_lockstep_packing_plan( + batch, + cfg, + batch_uid=9, + data_parallel_size=1, + ) + assert plan is not None + + log_csd_packing_telemetry(plan, batch) + + output = capsys.readouterr().out + assert "logical_samples=2 physical_bins=1 multi_sample_bins=1" in output + assert "side=student" in output + assert "side=teacher_0" in output + + +def test_logical_batch_digest_is_deterministic_and_content_bound(capsys): + batch = _make_batch(batch_size=2, t_student=4, t_teacher=6) + batch["kd_token_mask"] = batch["token_mask"].clone() + batch["student_semantic_regions"] = [ + ((0, "assistant", "content", 1, 3),), + ((0, "assistant", "eot", 3, 4),), + ] + batch["teacher_0_semantic_regions"] = [ + ((0, "assistant", "content", 2, 5),), + ((0, "assistant", "eot", 5, 6),), + ] + + first = build_csd_logical_batch_digest_record(batch, batch_uid=9) + second = build_csd_logical_batch_digest_record(batch, batch_uid=9) + assert first == second + assert first["logical_samples"] == 2 + assert first["sample_ids"]["count"] == 2 + assert first["fields"]["input_ids"]["shape"] == [2, 4] + assert first["fields"]["teacher_0_input_ids"]["shape"] == [2, 6] + assert "student_semantic_regions" in first["fields"] + assert "alignment_0_pair_valid" in first["fields"] + + batch["input_ids"][0, 0] = 1 + changed = build_csd_logical_batch_digest_record(batch, batch_uid=9) + assert ( + changed["fields"]["input_ids"]["sha256"] + != first["fields"]["input_ids"]["sha256"] + ) + assert changed["record_sha256"] != first["record_sha256"] + + logged = log_csd_logical_batch_digest(batch, batch_uid=9) + output = capsys.readouterr().out + assert logged == changed + assert output.startswith("CSD_LOGICAL_BATCH_DIGEST {") + assert "sample-0" not in output + + +def test_dense_ipc_telemetry_counts_logical_shard_bytes(): + handles = [ + { + "teacher_shards": [ + {"actual_shape": (4, 8), "dtype": torch.float32}, + {"actual_shape": (4, 8), "dtype": torch.float32}, + ] + }, + { + "teacher_shards": [ + {"actual_shape": (2, 8), "dtype": torch.bfloat16}, + ] + }, + ] + + assert csd_mod._dense_ipc_logical_bytes(handles) == 288 + + +# --------------------------------------------------------------------------- +# averaged_logits direct-KL guard +# --------------------------------------------------------------------------- + + +def _bare_averaged_logits_loss_fn(projection_matrix_paths): + """A ``CSDLossFn`` carrying only the attrs the + ``averaged_logits`` guard reads, built via ``__new__`` to skip the heavy + config-driven ``__init__``.""" + fn = CSDLossFn.__new__(CSDLossFn) + fn.num_teachers = len(projection_matrix_paths) + fn.projection_matrix_paths = list(projection_matrix_paths) + fn.student_to_teacher_table_paths = [None] * fn.num_teachers + fn.teacher_weights = [1.0] * fn.num_teachers + return fn + + +def test_averaged_logits_cross_tokenizer_skips_direct_kl_fast_path(monkeypatch): + # Two cross-tokenizer teachers (non-null projection paths) whose logits + # happen to share a shape must NOT take the direct per-position KL fast + # path: it assumes the student's tokenizer and would mismatch the + # student's token_mask when the teacher length differs (T_t != T_s). + fn = _bare_averaged_logits_loss_fn(["t0_proj.pt", "t1_proj.pt"]) + fn.teacher_weights = [2.0, 3.0] # distinct weights so weighting is exercised + + fallback_calls = [] + + def _fallback(i, *args, **kwargs): + fallback_calls.append(i) + # Distinct nonzero per-teacher KD so the weighted sum is non-trivial: + # a "last teacher overwrites" or "weight ignored" bug would change it. + return torch.tensor(float(i + 1)), {} + + def _fast_path(*args, **kwargs): + raise AssertionError( + "same_tokenizer_full_vocab_kl reached for cross-tokenizer teachers" + ) + + fn._compute_teacher_kd = _fallback + monkeypatch.setattr( + "nemo_rl.algorithms.loss.csd.same_tokenizer_full_vocab_kl", + _fast_path, + ) + + # Equal-shaped teacher logits (the misleading signal) but teacher length 13 + # vs the student's length 10. + teacher_logits = torch.zeros(2, 13, 24) + teacher_full = {0: teacher_logits, 1: teacher_logits.clone()} + student_logits = torch.zeros(2, 10, 32) + + total_kd, metrics = fn._averaged_logits_kd( + student_logits, + {}, + teacher_full, + {}, + torch.tensor(20.0), + teacher_sparse_logits_by_idx={}, + tp_group=None, + cp_group=None, + ) + assert fallback_calls == [0, 1] # per-teacher fallback path, one call each + # total_kd = Σ_i weight_i * kd_i = 2*1 + 3*2 = 8. + assert total_kd.item() == pytest.approx(2.0 * 1.0 + 3.0 * 2.0) + assert metrics["teacher_0/weighted_kl"] == pytest.approx(2.0) + assert metrics["teacher_1/weighted_kl"] == pytest.approx(6.0) + + +def test_averaged_logits_same_tokenizer_takes_direct_kl_fast_path(monkeypatch): + # All-null projection paths => genuine same-tokenizer teachers, so the + # averaging + single direct-KL fast path is the correct branch. + fn = _bare_averaged_logits_loss_fn([None, None]) + fn.temperature = 1.0 + fn.reverse_kl = False + fn.student_vocab_size = 30 + + fast_calls = [] + + def _fast_path(*args, **kwargs): + fast_calls.append(kwargs["vocab_size"]) + return torch.tensor(1.23) + + monkeypatch.setattr( + "nemo_rl.algorithms.loss.csd.same_tokenizer_full_vocab_kl", + _fast_path, + ) + fn._compute_teacher_kd = MagicMock( + side_effect=AssertionError("fallback reached for same-tokenizer teachers") + ) + + logits = torch.zeros(2, 10, 32) + teacher_full = {0: logits.clone(), 1: logits.clone()} + + kd, metrics = fn._averaged_logits_kd( + logits, + {}, + teacher_full, + {0: MagicMock()}, + torch.tensor(20.0), + teacher_sparse_logits_by_idx={}, + tp_group=None, + cp_group=None, + ) + assert fast_calls == [30] # fast path taken once, cut to the tokenizer vocab + assert "kl_loss" in metrics + assert kd.item() == pytest.approx(1.23) + assert metrics["teacher_0/weighted_kl"] == pytest.approx(1.23 / 2) + assert metrics["teacher_1/weighted_kl"] == pytest.approx(1.23 / 2) + + +# --------------------------------------------------------------------------- +# teacher weight-metric score masking +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("metric", ["ce", "entropy", "max_prob"]) +def test_teacher_weight_score_ignores_padded_positions(metric): + # The per-teacher weight/selection score must exclude padded positions: + # padding logits are near-uniform noise and would otherwise dominate the + # score on long-padded batches. Corrupting everything in the padded tail + # must not change the score. + fn = CSDLossFn.__new__(CSDLossFn) + fn.sum_weights_metric = metric + + torch.manual_seed(0) + batch, seqlen, vocab = 2, 5, 8 + logits = torch.randn(batch, seqlen, vocab) + ids = torch.randint(0, vocab, (batch, seqlen)) + # Last two positions of each sample are padding. + token_mask = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 0, 0]], dtype=torch.float32) + sample_mask = torch.ones(batch) + + score = fn._teacher_weight_score(logits, ids, token_mask, sample_mask) + + # Corrupt the padded tail; a properly masked score must not move. + logits_corrupt = logits.clone() + logits_corrupt[:, 3:, :] = 1e4 + ids_corrupt = ids.clone() + ids_corrupt[:, 3:] = 0 + score_corrupt = fn._teacher_weight_score( + logits_corrupt, ids_corrupt, token_mask, sample_mask + ) + + assert torch.allclose(score, score_corrupt, atol=1e-5) + + +def test_teacher_weight_score_masks_dropped_samples(): + # sample_mask=0 rows must not contribute either. + fn = CSDLossFn.__new__(CSDLossFn) + fn.sum_weights_metric = "max_prob" + + torch.manual_seed(0) + logits = torch.randn(2, 4, 8) + ids = torch.randint(0, 8, (2, 4)) + token_mask = torch.ones(2, 4) + sample_mask = torch.tensor([1.0, 0.0]) # second sample dropped + + score = fn._teacher_weight_score(logits, ids, token_mask, sample_mask) + + logits_corrupt = logits.clone() + logits_corrupt[1] = 1e4 # corrupt the dropped sample + score_corrupt = fn._teacher_weight_score( + logits_corrupt, ids, token_mask, sample_mask + ) + + assert torch.allclose(score, score_corrupt, atol=1e-5) + + +# --------------------------------------------------------------------------- +# sum_weights_metric is a sum-mode-only knob +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("mode", ["averaged_logits", "select_teacher"]) +def test_sum_weights_metric_rejected_outside_sum_mode(mode): + # Dynamic weighting is only applied in kd_loss_mode="sum" (matching the + # reference). Combining sum_weights_metric with another mode used to be + # silently dropped; it must now fail loudly at construction. + with pytest.raises(ValueError, match="sum_weights_metric"): + CSDLossFn( + { + "kd_loss_mode": mode, + "sum_weights_metric": "ce", + } + ) + + +@pytest.mark.parametrize("weights", [[0.0, 0.0], [1.0, -1.0]]) +def test_averaged_logits_rejects_zero_weight_sum(weights): + # averaged_logits divides by sum(teacher_weights) to form a convex average, + # so a zero weight-sum (all zeros, or a signed set cancelling to 0) is + # rejected at construction with a clear message instead of failing with a + # deep ZeroDivisionError mid-step. + with pytest.raises(ValueError, match="must not sum to zero"): + CSDLossFn( + { + "kd_loss_mode": "averaged_logits", + "teacher_weights": weights, + } + ) + + +# --------------------------------------------------------------------------- +# DP-global teacher score (cross-rank agreement primitive) +# --------------------------------------------------------------------------- + + +def test_dp_global_masked_mean_single_process_is_local_masked_mean(): + # Without a process group the DP all-reduce is a no-op, so the global + # masked mean reduces to the plain masked mean: padded positions excluded, + # and the result is detached (it gates selection / weighting, never + # back-propagated). Cross-rank agreement across DP ranks is verified by + # ``test_dp_global_masked_mean_agrees_across_ranks`` below. + fn = CSDLossFn.__new__(CSDLossFn) + values = torch.tensor([[1.0, 2.0, 100.0]], requires_grad=True) + mask = torch.tensor([[1.0, 1.0, 0.0]]) + + out = fn._dp_global_masked_mean(values, mask) + + assert out.item() == pytest.approx(1.5) # (1+2)/2; padded 100.0 excluded + assert not out.requires_grad # detached + + +def _dp_global_masked_mean_worker(rank, world_size, init_file, q): + os.environ["MASTER_ADDR"] = "127.0.0.1" + dist.init_process_group( + backend="gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + ) + try: + fn = CSDLossFn.__new__(CSDLossFn) + # Deliberately uneven shards: rank 1's masked-out 1e9 must not leak, + # and the mean must combine across ranks (not stay rank-local). + if rank == 0: + values, mask = torch.tensor([[1.0, 2.0]]), torch.tensor([[1.0, 1.0]]) + else: + values, mask = torch.tensor([[10.0, 1e9]]), torch.tensor([[1.0, 0.0]]) + q.put((rank, fn._dp_global_masked_mean(values, mask).item())) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not has_gloo(), reason="gloo backend unavailable") +def test_dp_global_masked_mean_agrees_across_ranks(tmp_path): + """Both DP ranks compute the identical DP-global masked mean = 13/3. + + rank 0 contributes (1+2) over 2 valid positions, rank 1 contributes 10 over + 1 (its 1e9 is masked out); the all-reduced sum/count give (3+10)/(2+1), + proving the score is DP-global (not rank-local) and padding is excluded. + """ + ctx = mp.get_context("spawn") + q = ctx.Queue() + procs = [ + ctx.Process( + target=_dp_global_masked_mean_worker, + args=(rank, 2, str(tmp_path / "init"), q), + ) + for rank in range(2) + ] + for p in procs: + p.start() + for p in procs: + p.join(timeout=60) + assert p.exitcode == 0 + results = dict(q.get() for _ in range(len(procs))) + assert results[0] == results[1] + assert results[0] == pytest.approx(13.0 / 3.0) + + +# --------------------------------------------------------------------------- +# select_teacher chooses the lowest-CE teacher +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("better", [0, 1]) +def test_select_teacher_picks_lowest_ce_teacher(better): + # select_teacher must use the teacher with the lowest next-token CE on its + # own tokens. Two same-vocab teachers: the "better" one peaks on the true + # next token (CE ~ 0); the other is uniform (CE = log V). Parametrized over + # which index is better so it can't pass by always returning 0. + fn = CSDLossFn.__new__(CSDLossFn) + fn.num_teachers = 2 + fn.student_to_teacher_table_paths = [None, None] + fn.projection_matrix_paths = [None, None] # same-vocab: scored on student ids + + vocab, seqlen = 8, 4 + input_ids = torch.tensor([[1, 2, 3, 0]]) + data = { + "input_ids": input_ids, + "token_mask": torch.ones(1, seqlen), + "sample_mask": torch.ones(1), + } + + # Same-vocab teachers are scored on the student tokens via the thin + # alignment (``align.student_input_ids`` / ``align.student_token_mask``). + align = SimpleNamespace( + student_input_ids=input_ids, + student_kd_token_mask=None, + student_token_mask=torch.ones(1, seqlen), + ) + aligns_by_idx = {0: align, 1: align} + + confident = torch.full((1, seqlen, vocab), -10.0) + for t in range(seqlen - 1): + confident[0, t, input_ids[0, t + 1]] = 10.0 # peak on the true next token + uniform = torch.zeros(1, seqlen, vocab) + teacher_full = {better: confident, 1 - better: uniform} + + selected = [] + + def _fake_kd(i, *args, **kwargs): + selected.append(i) + return torch.tensor(2.5), {"kl_loss": 2.5} + + fn._compute_teacher_kd = _fake_kd + + _kd, metrics = fn._select_teacher_kd( + torch.zeros(1, seqlen, vocab), + data, + teacher_full, + aligns_by_idx, + torch.tensor(3.0), + teacher_sparse_logits_by_idx={}, + tp_group=None, + cp_group=None, + ) + + assert metrics["selected_teacher"] == better + assert selected == [better] # KD computed only for the selected teacher + assert metrics[f"teacher_{better}/weighted_kl"] == pytest.approx(2.5) + assert metrics[f"teacher_{1 - better}/weighted_kl"] == pytest.approx(0.0) + + +# --------------------------------------------------------------------------- +# sum-mode aggregation: weighted sum + normalize_teacher_by_vocab rescaling +# --------------------------------------------------------------------------- + + +def _bare_sum_kd_loss_fn( + *, + num_teachers, + teacher_weights, + normalize_teacher_by_vocab=False, + teacher_vocab_sizes=None, +): + """A ``CSDLossFn`` carrying only the attrs ``_sum_kd`` + reads, built via ``__new__`` to skip the config-driven ``__init__``. + Static weights only (``sum_weights_metric=None``).""" + fn = CSDLossFn.__new__(CSDLossFn) + fn.num_teachers = num_teachers + fn.teacher_weights = list(teacher_weights) + fn.sum_weights_metric = None + fn.normalize_teacher_by_vocab = normalize_teacher_by_vocab + fn.teacher_vocab_sizes = ( + list(teacher_vocab_sizes) if teacher_vocab_sizes else [1] * num_teachers + ) + return fn + + +def test_sum_kd_weighted_sum_of_per_teacher_kd(): + # total_kd = Σ_i weight_i * KD_i. Distinct nonzero per-teacher KD + distinct + # weights so "last teacher overwrites" or "per-teacher weight ignored" can't + # pass silently. + fn = _bare_sum_kd_loss_fn(num_teachers=2, teacher_weights=[2.0, 3.0]) + + def _fake_kd(i, *args, **kwargs): + return torch.tensor(float(i + 1)), {"kl_loss": float(i + 1)} + + fn._compute_teacher_kd = _fake_kd + + total_kd, per_metrics = fn._sum_kd( + torch.zeros(1, 4, 8), # student_logits_contig (only used for device/dtype) + {}, + {}, + {}, + torch.tensor(10.0), + teacher_sparse_logits_by_idx={}, + tp_group=None, + cp_group=None, + ) + # 2*1 + 3*2 = 8. + assert total_kd.item() == pytest.approx(2.0 * 1.0 + 3.0 * 2.0) + # Per-teacher KD metrics are suffixed and the static weights are reported. + assert per_metrics["kl_loss_t0"] == 1.0 + assert per_metrics["kl_loss_t1"] == 2.0 + assert per_metrics["weight_t0"] == 2.0 + assert per_metrics["weight_t1"] == 3.0 + assert per_metrics["teacher_0/weighted_kl"] == pytest.approx(2.0) + assert per_metrics["teacher_1/weighted_kl"] == pytest.approx(6.0) + + +def test_sum_kd_normalize_teacher_by_vocab_rescales_by_log_ratio(): + # With normalize_teacher_by_vocab, each teacher's KD is scaled by + # log(V_t_i) / log(min_j V_t_j). V0 = min so its scale is exactly 1; V1 = V0^2 + # so its scale = 2. A min<->max or numerator/denominator swap changes this. + v0, v1 = 100, 10000 # log(v1)/log(v0) == 2 + fn = _bare_sum_kd_loss_fn( + num_teachers=2, + teacher_weights=[1.0, 1.0], + normalize_teacher_by_vocab=True, + teacher_vocab_sizes=[v0, v1], + ) + + def _fake_kd(i, *args, **kwargs): + return torch.tensor(float(i + 1)), {} + + fn._compute_teacher_kd = _fake_kd + + total_kd, per_metrics = fn._sum_kd( + torch.zeros(1, 4, 8), + {}, + {}, + {}, + torch.tensor(10.0), + teacher_sparse_logits_by_idx={}, + tp_group=None, + cp_group=None, + ) + s0 = math.log(v0) / math.log(v0) # 1.0 + s1 = math.log(v1) / math.log(v0) # 2.0 + # kd = [1, 2], weights = [1, 1]: 1*1*s0 + 2*1*s1. + assert total_kd.item() == pytest.approx(1.0 * s0 + 2.0 * s1) + assert per_metrics["teacher_0/weighted_kl"] == pytest.approx(1.0 * s0) + assert per_metrics["teacher_1/weighted_kl"] == pytest.approx(2.0 * s1) + + +def test_teacher_routing_metrics_use_each_teachers_tokenization(): + """Routing counts exclude padded rows and use each teacher's token axis.""" + fn = CSDLossFn.__new__(CSDLossFn) + fn.num_teachers = 2 + fn.student_to_teacher_table_paths = [None, None] + fn.projection_matrix_paths = ["/tmp/projection.pt", None] + data = BatchedDataDict( + { + "sample_mask": torch.tensor([1, 0]), + "token_mask": torch.tensor([[1, 1, 1, 0], [1, 1, 1, 1]]), + "teacher_0_token_mask": torch.tensor([[1, 1, 1, 1, 0], [1, 1, 1, 1, 1]]), + } + ) + + metrics = fn._teacher_routing_metrics(data) + + assert metrics == { + "teacher_0/routed_samples": 1, + "teacher_0/routed_tokens": 4, + "teacher_1/routed_samples": 1, + "teacher_1/routed_tokens": 3, + } + + +# --------------------------------------------------------------------------- +# dynamic teacher weights: softmax(alpha * per-teacher scores) across teachers +# --------------------------------------------------------------------------- + + +def _dynamic_weights_fn(alpha, scores): + """``CSDLossFn`` stubbed to feed fixed per-teacher + scores into ``_compute_dynamic_weights`` (isolates the aggregation).""" + fn = CSDLossFn.__new__(CSDLossFn) + fn.num_teachers = len(scores) + fn.alpha = alpha + fn.normalize_teacher_by_vocab = False + fn._teacher_score_inputs = lambda i, *a, **k: (None, None, None) + it = iter([torch.tensor(float(s)) for s in scores]) + fn._teacher_weight_score = lambda *a, **k: next(it) + return fn + + +def test_compute_dynamic_weights_softmax_over_teachers(): + # Weights are softmax(alpha * scores) over the teacher axis. Stubbed scores + # isolate the aggregation, catching "softmax over wrong axis" or "alpha + # applied to logits not scores". + data = {"input_ids": torch.zeros(1, 1), "sample_mask": torch.ones(1)} + + weights = _dynamic_weights_fn(1.0, [1.0, 0.0])._compute_dynamic_weights( + data, {}, {} + ) + expected = torch.softmax(torch.tensor([1.0, 0.0]), dim=0) + assert len(weights) == 2 + assert weights[0].item() == pytest.approx(expected[0].item()) + assert weights[1].item() == pytest.approx(expected[1].item()) + assert (weights[0] + weights[1]).item() == pytest.approx(1.0) + assert weights[0].item() > weights[1].item() # higher score -> higher weight + + # alpha multiplies the scores before softmax: a larger alpha sharpens onto + # the top-scoring teacher. + sharp = _dynamic_weights_fn(8.0, [1.0, 0.0])._compute_dynamic_weights(data, {}, {}) + assert sharp[0].item() == pytest.approx( + torch.softmax(torch.tensor([8.0, 0.0]), dim=0)[0].item() + ) + assert sharp[0].item() > weights[0].item() diff --git a/tests/unit/algorithms/csd/test_global_chunk_normalization.py b/tests/unit/algorithms/csd/test_global_chunk_normalization.py new file mode 100644 index 00000000000..c5aaba1e2a5 --- /dev/null +++ b/tests/unit/algorithms/csd/test_global_chunk_normalization.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock + +import pytest +import torch + +from nemo_rl.algorithms.loss.csd import CSDLossFn + + +def test_cross_token_teacher_receives_its_global_chunk_count(): + loss_fn = object.__new__(CSDLossFn) + loss_fn.projection_matrix_paths = [None, None, "projection.pt"] + loss_fn.teacher_vocab_sizes = [16, 16, 32] + loss_fn._compute_csd_loss = MagicMock(return_value=(torch.tensor(1.5), {})) + counts = {0: torch.tensor(11), 2: torch.tensor(29)} + + loss_fn._compute_teacher_kd( + 2, + torch.zeros(1), + MagicMock(), + {2: torch.zeros(1)}, + {2: MagicMock()}, + torch.tensor(7), + teacher_sparse_logits_by_idx={}, + tp_group=None, + cp_group=None, + global_valid_chunks_by_idx=counts, + ) + + assert ( + loss_fn._compute_csd_loss.call_args.kwargs["global_valid_chunks"] is counts[2] + ) + + +def test_missing_cross_token_teacher_count_does_not_fall_back_to_microbatch(): + loss_fn = object.__new__(CSDLossFn) + loss_fn.projection_matrix_paths = ["projection.pt"] + loss_fn.teacher_vocab_sizes = [32] + + with pytest.raises(KeyError): + loss_fn._compute_teacher_kd( + 0, + torch.zeros(1), + MagicMock(), + {0: torch.zeros(1)}, + {0: MagicMock()}, + torch.tensor(7), + teacher_sparse_logits_by_idx={}, + tp_group=None, + cp_group=None, + global_valid_chunks_by_idx={}, + ) diff --git a/tests/unit/algorithms/csd/test_packing_loss.py b/tests/unit/algorithms/csd/test_packing_loss.py new file mode 100644 index 00000000000..abd39fd7edc --- /dev/null +++ b/tests/unit/algorithms/csd/test_packing_loss.py @@ -0,0 +1,411 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from nemo_rl.algorithms.loss.csd_packing import ( + CSD_LOGICAL_METRICS_KEY, + CSDSequencePackingLossWrapper, + restore_packed_student_logits, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +def test_restore_packed_student_logits_uses_padded_offsets_and_raw_lengths(): + packed = torch.tensor( + [ + [ + [1.0], + [2.0], + [3.0], + [99.0], + [10.0], + [11.0], + [12.0], + [13.0], + [14.0], + [98.0], + [97.0], + [96.0], + ] + ] + ) + raw_cu = torch.tensor([0, 3, 8], dtype=torch.int32) + padded_cu = torch.tensor([0, 4, 12], dtype=torch.int32) + + first = restore_packed_student_logits( + packed, + sequence_index=0, + cu_seqlens_q=raw_cu, + cu_seqlens_q_padded=padded_cu, + ) + second = restore_packed_student_logits( + packed, + sequence_index=1, + cu_seqlens_q=raw_cu, + cu_seqlens_q_padded=padded_cu, + ) + + assert first.shape == (1, 4, 1) + assert second.shape == (1, 8, 1) + assert first[:, :, 0].tolist() == [[1.0, 2.0, 3.0, 0.0]] + assert second[:, :, 0].tolist() == [[10.0, 11.0, 12.0, 13.0, 14.0, 0.0, 0.0, 0.0]] + + +@pytest.mark.parametrize( + ("cp_rank", "expected"), + [ + (0, [[10.0, 11.0, 0.0, 0.0]]), + (1, [[12.0, 13.0, 14.0, 0.0]]), + ], +) +def test_restore_packed_student_logits_cp2_reapplies_logical_head_tail_layout( + cp_rank, expected +): + import nemo_rl.algorithms.loss.csd_packing as packing_loss + + cp_group = MagicMock() + packed_local = torch.arange(6, dtype=torch.float32).reshape(1, 6, 1) + gathered_physical = torch.tensor( + [[[10.0], [11.0], [12.0], [13.0], [14.0], [99.0], [98.0], [97.0]]] + ) + with ( + patch.object(torch.distributed, "get_world_size", return_value=2), + patch.object(torch.distributed, "get_rank", return_value=cp_rank), + patch.object( + packing_loss, + "allgather_cp_sharded_tensor", + return_value=gathered_physical, + ) as mock_allgather, + ): + restored = restore_packed_student_logits( + packed_local, + sequence_index=1, + cu_seqlens_q=torch.tensor([0, 3, 8], dtype=torch.int32), + cu_seqlens_q_padded=torch.tensor([0, 4, 12], dtype=torch.int32), + context_parallel_group=cp_group, + ) + + assert restored[:, :, 0].tolist() == expected + # Sequence 1 begins at padded offset 4 globally, i.e. local offset 2. + assert mock_allgather.call_args.args[0].shape == (1, 4, 1) + + +def test_csd_packing_wrapper_preserves_asymmetric_teacher_and_pair_axes(): + data = BatchedDataDict( + { + "input_ids": torch.arange(16).reshape(2, 8), + "input_lengths": torch.tensor([3, 5]), + "token_mask": torch.tensor( + [ + [1, 1, 1, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 0, 0, 0], + ] + ), + "kd_token_mask": torch.tensor( + [ + [0, 1, 1, 0, 0, 0, 0, 0], + [1, 0, 1, 0, 1, 0, 0, 0], + ] + ), + "sample_mask": torch.ones(2), + # Deliberately asymmetric with T_student=8. + "teacher_0_input_ids": torch.arange(18).reshape(2, 9), + "teacher_0_token_mask": torch.ones(2, 9), + "alignment_0_teacher_chunk_id": torch.arange(18).reshape(2, 9), + "alignment_0_student_chunk_id": torch.arange(16).reshape(2, 8), + "alignment_0_student_exact_partition_mask": torch.ones( + 2, 8, dtype=torch.bool + ), + # Pair axis P=4 must not be truncated to either raw student length. + "alignment_0_pair_valid": torch.ones(2, 4, dtype=torch.bool), + "alignment_0_pair_is_correct": torch.ones(2, 4, dtype=torch.bool), + "student_semantic_regions": [[("student", 0)], [("student", 1)]], + "teacher_0_semantic_regions": [[("teacher", 0)], [("teacher", 1)]], + "teacher_0_full_logits_ipc": [{"row": 0}, {"row": 1}], + } + ) + original_tensors = { + key: value.clone() + for key, value in data.items() + if isinstance(value, torch.Tensor) + } + packed = torch.tensor( + [ + [ + [1.0], + [2.0], + [3.0], + [99.0], + [10.0], + [11.0], + [12.0], + [13.0], + [14.0], + [98.0], + [97.0], + [96.0], + ] + ], + requires_grad=True, + ) + observed: list[dict] = [] + + def prepare_fn(*, logits, data, **_kwargs): + observed.append( + { + "logits": logits.detach().clone(), + "shapes": { + key: tuple(value.shape) + for key, value in data.items() + if isinstance(value, torch.Tensor) + }, + "input_lengths": data["input_lengths"].clone(), + "input_ids": data["input_ids"].clone(), + "token_mask": data["token_mask"].clone(), + "kd_token_mask": data["kd_token_mask"].clone(), + "student_regions": data["student_semantic_regions"], + "teacher_regions": data["teacher_0_semantic_regions"], + "teacher_ipc": data["teacher_0_full_logits_ipc"], + } + ) + return {"logical_logits": logits}, data + + class RecordingLoss: + def __call__(self, *, logical_logits, **_kwargs): + return logical_logits.sum(), {"logical_calls": 1} + + wrapper = CSDSequencePackingLossWrapper( + loss_fn=RecordingLoss(), + prepare_fn=prepare_fn, + cu_seqlens_q=torch.tensor([0, 3, 8], dtype=torch.int32), + cu_seqlens_q_padded=torch.tensor([0, 4, 12], dtype=torch.int32), + ) + loss, metrics = wrapper( + packed, + data, + global_valid_seqs=torch.tensor(2.0), + global_valid_toks=torch.tensor(5.0), + ) + loss.backward() + + assert len(observed) == 2 + assert observed[0]["logits"][:, :, 0].tolist() == [[1.0, 2.0, 3.0, 0.0]] + assert observed[1]["logits"][:, :, 0].tolist() == [ + [10.0, 11.0, 12.0, 13.0, 14.0, 0.0, 0.0, 0.0] + ] + restored_positions = sum(item["logits"].shape[1] for item in observed) + padded_cu = wrapper.cu_seqlens_q_padded + expected_effective_positions = sum( + right - left for left, right in zip(padded_cu[:-1], padded_cu[1:]) + ) + assert restored_positions == expected_effective_positions == 12 + assert restored_positions != data.size * data["input_ids"].shape[1] + assert [item["shapes"]["input_ids"] for item in observed] == [(1, 4), (1, 8)] + assert [item["shapes"]["token_mask"] for item in observed] == [(1, 4), (1, 8)] + assert [item["shapes"]["kd_token_mask"] for item in observed] == [ + (1, 4), + (1, 8), + ] + assert [item["shapes"]["alignment_0_student_chunk_id"] for item in observed] == [ + (1, 4), + (1, 8), + ] + assert [ + item["shapes"]["alignment_0_student_exact_partition_mask"] for item in observed + ] == [(1, 4), (1, 8)] + assert [item["input_lengths"].tolist() for item in observed] == [[3], [5]] + assert observed[0]["input_ids"].tolist() == [[0, 1, 2, 3]] + assert observed[1]["input_ids"].tolist() == [[8, 9, 10, 11, 12, 13, 14, 15]] + assert observed[0]["token_mask"].tolist() == [[1, 1, 1, 0]] + assert observed[0]["kd_token_mask"].tolist() == [[0, 1, 1, 0]] + for item in observed: + shapes = item["shapes"] + assert shapes["teacher_0_input_ids"] == (1, 9) + assert shapes["teacher_0_token_mask"] == (1, 9) + assert shapes["alignment_0_teacher_chunk_id"] == (1, 9) + assert shapes["alignment_0_pair_valid"] == (1, 4) + assert shapes["alignment_0_pair_is_correct"] == (1, 4) + assert [item["student_regions"] for item in observed] == [ + [[("student", 0)]], + [[("student", 1)]], + ] + assert [item["teacher_regions"] for item in observed] == [ + [[("teacher", 0)]], + [[("teacher", 1)]], + ] + assert [item["teacher_ipc"] for item in observed] == [ + [{"row": 0}], + [{"row": 1}], + ] + for key, original in original_tensors.items(): + torch.testing.assert_close(data[key], original) + assert data["student_semantic_regions"] == [ + [("student", 0)], + [("student", 1)], + ] + assert data["teacher_0_semantic_regions"] == [ + [("teacher", 0)], + [("teacher", 1)], + ] + assert data["teacher_0_full_logits_ipc"] == [{"row": 0}, {"row": 1}] + assert metrics == { + "logical_calls": 2, + CSD_LOGICAL_METRICS_KEY: [ + {"logical_calls": 1}, + {"logical_calls": 1}, + ], + } + # Only real tokens receive gradient. Physical padding on either side of a + # sample boundary cannot act as another sample's next-token predictor. + assert packed.grad[:, :, 0].tolist() == [ + [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0] + ] + + +def test_csd_packed_effective_rows_match_unpacked_masked_loss_and_gradients(): + class MaskedLoss: + def __call__(self, *, logical_logits, data, **_kwargs): + mask = (data["token_mask"] * data["kd_token_mask"]).to(logical_logits.dtype) + return (logical_logits.squeeze(-1) * mask).sum(), { + "masked_tokens": int(mask.sum().item()) + } + + def prepare_fn(*, logits, data, **_kwargs): + return {"logical_logits": logits}, data + + data = BatchedDataDict( + { + "input_ids": torch.arange(16).reshape(2, 8), + "input_lengths": torch.tensor([3, 5]), + "token_mask": torch.tensor( + [ + [1, 1, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 0, 0, 0], + ] + ), + "kd_token_mask": torch.tensor( + [ + [1, 0, 1, 0, 0, 0, 0, 0], + [0, 1, 1, 0, 1, 0, 0, 0], + ] + ), + "sample_mask": torch.ones(2), + } + ) + packed_logits = torch.tensor( + [ + [ + [1.0], + [2.0], + [3.0], + [99.0], + [10.0], + [11.0], + [12.0], + [13.0], + [14.0], + [98.0], + [97.0], + [96.0], + ] + ], + requires_grad=True, + ) + wrapper = CSDSequencePackingLossWrapper( + loss_fn=MaskedLoss(), + prepare_fn=prepare_fn, + cu_seqlens_q=torch.tensor([0, 3, 8], dtype=torch.int32), + cu_seqlens_q_padded=torch.tensor([0, 4, 12], dtype=torch.int32), + ) + packed_loss, packed_metrics = wrapper( + packed_logits, + data, + global_valid_seqs=torch.tensor(2.0), + global_valid_toks=torch.tensor(8.0), + global_valid_kd_toks=torch.tensor(5.0), + ) + packed_loss.backward() + + unpacked_logits = [ + torch.tensor( + [[[1.0], [2.0], [3.0], [0.0], [0.0], [0.0], [0.0], [0.0]]], + requires_grad=True, + ), + torch.tensor( + [[[10.0], [11.0], [12.0], [13.0], [14.0], [0.0], [0.0], [0.0]]], + requires_grad=True, + ), + ] + unpacked_loss = torch.zeros(()) + for row_index, row_logits in enumerate(unpacked_logits): + row_data = data.slice(row_index, row_index + 1) + row_loss, _ = MaskedLoss()( + logical_logits=row_logits, + data=row_data, + ) + unpacked_loss = unpacked_loss + row_loss + unpacked_loss.backward() + + torch.testing.assert_close(packed_loss, unpacked_loss) + expected_packed_grad = torch.cat( + [ + unpacked_logits[0].grad[:, :4], + unpacked_logits[1].grad[:, :8], + ], + dim=1, + ) + torch.testing.assert_close(packed_logits.grad, expected_packed_grad) + assert packed_metrics["masked_tokens"] == 4 + + +def test_packed_metrics_preserve_unpacked_mbs1_cohorts_for_uneven_bins(): + class PerSampleLoss: + def __call__(self, *, logical_logits, data, **_kwargs): + sample_id = int(data["batch_item_id"][0].item()) + valid = int(data["sample_mask"][0].item()) + return logical_logits.sum() * valid, { + "loss": float(sample_id + 1), + "accuracy": float(sample_id) / 10, + "ipc_reconstruction_fallbacks": 1, + "ipc_reconstruction_fallbacks_t0": 1, + "num_valid_samples": valid, + } + + def prepare_fn(*, logits, data, **_kwargs): + return {"logical_logits": logits}, data + + def run_bin(sample_ids, sample_masks): + count = len(sample_ids) + data = BatchedDataDict( + { + "input_ids": torch.zeros(count, 2, dtype=torch.long), + "input_lengths": torch.full((count,), 2), + "sample_mask": torch.tensor(sample_masks), + "batch_item_id": torch.tensor(sample_ids), + } + ) + wrapper = CSDSequencePackingLossWrapper( + loss_fn=PerSampleLoss(), + prepare_fn=prepare_fn, + cu_seqlens_q=torch.arange(0, 2 * count + 1, 2, dtype=torch.int32), + ) + logits = torch.ones(1, 2 * count, 1) + _, metrics = wrapper(logits, data, torch.tensor(2.0), torch.tensor(4.0)) + return metrics[CSD_LOGICAL_METRICS_KEY] + + # Physical bins have cardinalities two and one; row 1 is masked but still + # executes above. Flattened records match three unpacked MBS1 calls with + # the same masked-row filtering, rather than two physical-bin means. + packed_records = run_bin([0, 1], [1, 0]) + run_bin([2], [1]) + unpacked_records = run_bin([0], [1]) + run_bin([1], [0]) + run_bin([2], [1]) + assert packed_records == unpacked_records + assert [record["accuracy"] for record in packed_records] == [0.0, 0.2] + # Reconstruction runs before the masked-row check, but the worker's public + # metric stream deliberately contains valid logical MBS1 cohorts only. + assert sum(record["ipc_reconstruction_fallbacks"] for record in packed_records) == 2 + assert ( + sum(record["ipc_reconstruction_fallbacks_t0"] for record in packed_records) == 2 + ) diff --git a/tests/unit/algorithms/test_loss_functions.py b/tests/unit/algorithms/test_loss_functions.py index 439890d9121..eabf7090b1a 100644 --- a/tests/unit/algorithms/test_loss_functions.py +++ b/tests/unit/algorithms/test_loss_functions.py @@ -25,6 +25,8 @@ DistillationLossFn, DPOLossConfig, DPOLossFn, + MseValueLossConfig, + MseValueLossFn, NLLLossFn, prepare_loss_input, ) @@ -3372,3 +3374,58 @@ def test_vocab_parallel_gather_columns_tp_sharded(monkeypatch): ref[..., idx].float().backward(grad_out) torch.testing.assert_close(shards[0].grad, ref.grad[..., :v_local]) torch.testing.assert_close(shards[1].grad, ref.grad[..., v_local:]) + + +def test_distillation_loss_num_valid_samples_respects_sample_mask(): + """The metric counts participating rows, not padded microbatch capacity.""" + torch.manual_seed(0) + student_log_probs = torch.log_softmax(torch.randn(2, 3, 4), dim=-1) + teacher_log_probs = torch.log_softmax(torch.randn(2, 3, 4), dim=-1) + sample_mask = torch.tensor([1.0, 0.0]) + data = { + "input_ids": torch.zeros(2, 4, dtype=torch.long), + "input_lengths": torch.full((2,), 4, dtype=torch.long), + "token_mask": torch.ones(2, 4), + "sample_mask": sample_mask, + "teacher_topk_logits": torch.empty(0), + "teacher_topk_indices": torch.empty(0, dtype=torch.long), + } + loss_fn = DistillationLossFn( + DistillationLossConfig( + kl_type="forward", + mixed_kl_weight=0.5, + zero_outside_topk=False, + ) + ) + + _, metrics = loss_fn( + student_log_probs, + teacher_log_probs, + None, + data, + global_valid_seqs=sample_mask.sum(), + global_valid_toks=torch.tensor(3.0), + ) + + assert metrics["num_valid_samples"] == 1 + + +def test_mse_value_loss_num_valid_samples_respects_sample_mask(): + loss_fn = MseValueLossFn(MseValueLossConfig()) + sample_mask = torch.tensor([1.0, 0.0]) + data = BatchedDataDict( + { + "token_mask": torch.ones(2, 3), + "sample_mask": sample_mask, + "returns": torch.zeros(2, 3), + } + ) + + _, metrics = loss_fn( + torch.zeros(2, 3), + data, + global_valid_seqs=sample_mask.sum(), + global_valid_toks=torch.tensor(3.0), + ) + + assert metrics["num_valid_samples"] == 1 diff --git a/tests/unit/algorithms/x_token/test_loss_utils.py b/tests/unit/algorithms/x_token/test_loss_utils.py index 598bb59916e..821e6c58e70 100644 --- a/tests/unit/algorithms/x_token/test_loss_utils.py +++ b/tests/unit/algorithms/x_token/test_loss_utils.py @@ -51,7 +51,10 @@ get_topk_projection, localize_alignment, parse_projection_file, + prepare_csd_cross_tokenizer_loss_input, prepare_xtoken_cross_tokenizer_loss_input, + rebuild_teacher_full_logits_from_ipc, + rebuild_teacher_sparse_logits_from_ipc, same_tokenizer_full_vocab_kl, same_tokenizer_topk_kl, select_teacher_topk_indices, @@ -72,6 +75,7 @@ ) from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.distributed.worker_groups import RayWorkerBuilder, RayWorkerGroup +from nemo_rl.models.policy.utils import DENSE_TEACHER_IPC_FLAT_LAYOUT def test_automodel_cp_layout_localizes_xtoken_windows_after_global_shift(): @@ -101,7 +105,7 @@ def test_automodel_cp_layout_localizes_xtoken_windows_after_global_shift(): patch( "nemo_rl.algorithms.x_token.loss_utils." "rebuild_teacher_full_logits_from_ipc", - return_value=teacher_logits, + return_value=(teacher_logits, 0), ) as rebuild_teacher, ): student_logits, teachers, aligns, tp_group, returned_cp_group = ( @@ -158,6 +162,102 @@ def test_automodel_cp_layout_rejects_non_divisible_student_sequence(): ) +def test_csd_automodel_cp_layout_localizes_csd_windows_after_global_shift(): + """Automodel layout restoration precedes contiguous x-token localization.""" + cp_group = object() + local_logits = torch.randn(1, 3, 2) + full_student_logits = torch.arange(12, dtype=torch.float32).reshape(1, 6, 2) + teacher_logits = torch.randn(1, 2, 4) + cp_sharder = MagicMock() + cp_sharder.gather_token_tensor.return_value = full_student_logits + teacher_ipc = [{"opaque": "handle"}] + data = { + "input_ids": torch.arange(6).unsqueeze(0), + "token_mask": torch.ones(1, 6), + "sample_mask": torch.ones(1), + "teacher_0_full_logits_ipc": teacher_ipc, + "teacher_0_input_ids": torch.arange(20, 24).unsqueeze(0), + "alignment_0_student_chunk_id": torch.tensor([[10, 11, 12, 13, 14, 15]]), + "alignment_0_teacher_chunk_id": torch.tensor([[20, 21, 22, 23]]), + "alignment_0_pair_valid": torch.ones(1, 2, dtype=torch.bool), + "alignment_0_pair_is_correct": torch.ones(1, 2, dtype=torch.bool), + "alignment_0_num_chunks": torch.tensor([2]), + } + + with ( + patch("torch.cuda.current_device", return_value=0), + patch("torch.distributed.get_world_size", return_value=2), + patch("torch.distributed.get_rank", return_value=1), + patch( + "nemo_rl.algorithms.x_token.loss_utils." + "rebuild_teacher_full_logits_from_ipc", + return_value=(teacher_logits, 0), + ) as rebuild_teacher, + ): + ( + student_logits, + teachers, + sparse_teachers, + aligns, + dense_reconstruction_fallbacks, + tp_group, + returned_cp_group, + dp_cp_group, + ) = prepare_csd_cross_tokenizer_loss_input( + local_logits, + data, + projection_matrix_paths=["projection.pt"], + context_parallel_group=cp_group, + cp_sharder=cp_sharder, + ) + + torch.testing.assert_close(student_logits, full_student_logits[:, 3:6]) + assert teachers[0] is teacher_logits + assert sparse_teachers == {} + assert dense_reconstruction_fallbacks == {0: 0} + assert tp_group is None + assert returned_cp_group is cp_group + assert dp_cp_group is None + torch.testing.assert_close(aligns[0].student_input_ids, data["input_ids"][:, 3:6]) + torch.testing.assert_close(aligns[0].student_token_mask, data["token_mask"][:, 3:6]) + torch.testing.assert_close(aligns[0].student_chunk_id, torch.tensor([[14, 15, -1]])) + torch.testing.assert_close(aligns[0].teacher_chunk_id, torch.tensor([[23, -1]])) + cp_sharder.gather_token_tensor.assert_called_once_with( + local_logits, seq_dim=1, trim=True + ) + rebuild_teacher.assert_called_once_with(teacher_ipc, cp_group=cp_group, device=0) + + +def test_csd_automodel_cp_layout_rejects_non_divisible_student_sequence(): + cp_group = object() + local_logits = torch.randn(1, 3, 2) + cp_sharder = MagicMock() + cp_sharder.gather_token_tensor.return_value = torch.randn(1, 5, 2) + + with ( + patch("torch.cuda.current_device", return_value=0), + patch("torch.distributed.get_world_size", return_value=2), + pytest.raises( + ValueError, + match=( + "CSD student sequence length must be divisible by the student " + "context parallel size" + ), + ), + ): + prepare_csd_cross_tokenizer_loss_input( + local_logits, + {}, + projection_matrix_paths=[], + context_parallel_group=cp_group, + cp_sharder=cp_sharder, + ) + + cp_sharder.gather_token_tensor.assert_called_once_with( + local_logits, seq_dim=1, trim=True + ) + + # --------------------------------------------------------------------------- # chunk_average_log_probs # --------------------------------------------------------------------------- @@ -474,12 +574,12 @@ def test_relaxed_mode_collision_resolution(self, tmp_path: Path): assert out["uncommon_teacher"].tolist() == [3] def test_strict_mode_collision_picks_lowest_student(self, tmp_path: Path): - """Strict (xtoken_loss=False) exact maps require weight 1.0 and a - ``-1`` sentinel in the second slot. On a teacher collision the lowest - student index wins; fuzzy rows (second slot != -1) are excluded. + """Strict (xtoken_loss=False) exact maps require a slot-0 weight of + 1.0. On a teacher collision the lowest student index wins; fuzzy rows + (top weight below 1.0) are excluded. s0, s1 both exact-map to teacher 1 -> s0 wins (lowest index) - s2 fuzzy (second slot 0) -> excluded + s2 fuzzy (0.7 / 0.3) -> excluded s3 exact-maps to teacher 0 (unique) """ indices = [[1, -1], [1, -1], [2, 0], [0, -1]] @@ -495,6 +595,26 @@ def test_strict_mode_collision_picks_lowest_student(self, tmp_path: Path): assert out["uncommon_student"].tolist() == [1, 2] assert out["uncommon_teacher"].tolist() == [2] + def test_strict_mode_accepts_zero_weight_trailing_ids(self, tmp_path: Path): + """Strict exact maps key on a top weight of 1.0 alone. + + s0 -> teacher 1 @ 1.0 with a zero-weight trailing id (no -1 sentinel) + s1 fuzzy (0.7 / 0.3) -> uncommon + s2 -> teacher 0 @ 1.0 with the legacy -1 sentinel + """ + indices = [[1, 0], [2, 1], [0, -1]] + likelihoods = [[1.0, 0.0], [0.7, 0.3], [1.0, 0.0]] + out = build_exact_token_map( + self._write_topk(tmp_path, indices, likelihoods), + torch.device("cpu"), + xtoken_loss=False, + teacher_vocab_size=3, + ) + assert out["common_student"].tolist() == [0, 2] + assert out["common_teacher"].tolist() == [1, 0] + assert out["uncommon_student"].tolist() == [1] + assert out["uncommon_teacher"].tolist() == [2] + # NOTE: the TP/CP > 1 paths of vocab_parallel_log_softmax / vocab_parallel_full_log_softmax / # project_student_to_teacher_vocab / vocab_parallel_argmax / @@ -581,6 +701,42 @@ def _entry(sample_idx, vend=5, buf=0, payload=("p",)): ] } + @staticmethod + def _compact_entry( + *, + token_offset: int, + valid_seq_len: int, + stored_seq_len: int, + used_tokens: int, + capacity_tokens: int, + full_seq_len: int = 4, + vocab_size: int = 5, + payload="compact", + ): + return { + "teacher_shards": [ + { + "payload_ipc": payload, + "buf_idx": 0, + "sample_index_in_buf": 0, + "ipc_layout": DENSE_TEACHER_IPC_FLAT_LAYOUT, + "storage_shape": (capacity_tokens, vocab_size), + "storage_token_offset": token_offset, + "storage_used_tokens": used_tokens, + "storage_capacity_tokens": capacity_tokens, + "stored_seq_len": stored_seq_len, + "valid_seq_len": valid_seq_len, + "vocab_start_index": 0, + "vocab_end_index": vocab_size, + "global_seq_start": 0, + "actual_shape": (full_seq_len, vocab_size), + "full_seq_len": full_seq_len, + "full_vocab_size": vocab_size, + "dtype": torch.float32, + } + ] + } + def test_fastpath_returns_zero_copy_view(self, monkeypatch): storage = torch.arange(2 * 4 * 5, dtype=torch.float32).reshape(1, 2, 4, 5) monkeypatch.setattr( @@ -640,6 +796,239 @@ def test_none_when_samples_not_contiguous(self): is None ) + def test_rebuild_reports_zero_for_actual_zero_copy(self, monkeypatch): + storage = torch.arange(2 * 4 * 5, dtype=torch.float32).reshape(1, 2, 4, 5) + monkeypatch.setattr( + "nemo_rl.models.policy.utils.rebuild_cuda_tensor_from_ipc", + lambda _handle, _device: storage, + ) + + out, fallback_count = rebuild_teacher_full_logits_from_ipc( + [self._entry(0), self._entry(1)], cp_group=None, device=0 + ) + + assert fallback_count == 0 + assert out.data_ptr() == storage.data_ptr() + assert torch.equal(out, storage[0, :2]) + + def test_compact_mbs1_full_row_is_a_live_zero_copy_view(self, monkeypatch): + storage = torch.full((9, 5), -99.0) + intended_row = torch.arange(4 * 5, dtype=torch.float32).reshape(4, 5) + storage[3:7].copy_(intended_row) + monkeypatch.setattr( + "nemo_rl.models.policy.utils.rebuild_cuda_tensor_from_ipc", + lambda _handle, _device: storage, + ) + entry = self._compact_entry( + token_offset=3, + valid_seq_len=4, + stored_seq_len=4, + used_tokens=9, + capacity_tokens=9, + ) + + out, fallback_count = rebuild_teacher_full_logits_from_ipc( + [entry], cp_group=None, device="cpu" + ) + + assert fallback_count == 0 + assert out.data_ptr() == storage[3].data_ptr() + assert torch.equal(out[0], intended_row) + storage[4, 2] = -123.0 + assert out[0, 1, 2].item() == -123.0 + + def test_compact_short_row_reconstructs_exact_zero_suffix(self, monkeypatch): + storage = torch.arange(2 * 5, dtype=torch.float32).reshape(2, 5) + monkeypatch.setattr( + "nemo_rl.models.policy.utils.rebuild_cuda_tensor_from_ipc", + lambda _handle, _device: storage, + ) + entry = self._compact_entry( + token_offset=0, + valid_seq_len=2, + stored_seq_len=2, + used_tokens=2, + capacity_tokens=2, + ) + + out, fallback_count = rebuild_teacher_full_logits_from_ipc( + [entry], cp_group=None, device="cpu" + ) + + expected = torch.zeros(1, 4, 5) + expected[0, :2] = storage + assert fallback_count == 1 + assert torch.equal(out, expected) + + def test_compact_zero_row_never_maps_a_payload(self, monkeypatch): + monkeypatch.setattr( + "nemo_rl.models.policy.utils.rebuild_cuda_tensor_from_ipc", + lambda _handle, _device: pytest.fail("zero row must not map IPC"), + ) + entry = self._compact_entry( + token_offset=0, + valid_seq_len=0, + stored_seq_len=0, + used_tokens=0, + capacity_tokens=0, + payload=None, + ) + + out, fallback_count = rebuild_teacher_full_logits_from_ipc( + [entry], cp_group=None, device="cpu" + ) + + assert fallback_count == 1 + assert torch.equal(out, torch.zeros(1, 4, 5)) + + def test_compact_rebuilt_storage_shape_must_match_handle(self, monkeypatch): + storage = torch.zeros(3, 5) + monkeypatch.setattr( + "nemo_rl.models.policy.utils.rebuild_cuda_tensor_from_ipc", + lambda _handle, _device: storage, + ) + entry = self._compact_entry( + token_offset=0, + valid_seq_len=4, + stored_seq_len=4, + used_tokens=4, + capacity_tokens=4, + ) + + with pytest.raises(ValueError, match="unexpected storage shape"): + rebuild_teacher_full_logits_from_ipc([entry], cp_group=None, device="cpu") + + def test_rebuild_counts_rows_that_use_shard_assembly(self, monkeypatch): + vocab_left = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(1, 2, 4, 3) + vocab_right = 100 + torch.arange(2 * 4 * 2, dtype=torch.float32).reshape( + 1, 2, 4, 2 + ) + storages = {"left": vocab_left, "right": vocab_right} + monkeypatch.setattr( + "nemo_rl.models.policy.utils.rebuild_cuda_tensor_from_ipc", + lambda handle, _device: storages[handle], + ) + + entries = [] + for sample_idx in range(2): + entries.append( + { + "teacher_shards": [ + { + "payload_ipc": "left", + "buf_idx": 0, + "sample_index_in_buf": sample_idx, + "vocab_start_index": 0, + "vocab_end_index": 3, + "global_seq_start": 0, + "actual_shape": (4, 3), + "full_seq_len": 4, + "full_vocab_size": 5, + }, + { + "payload_ipc": "right", + "buf_idx": 0, + "sample_index_in_buf": sample_idx, + "vocab_start_index": 3, + "vocab_end_index": 5, + "global_seq_start": 0, + "actual_shape": (4, 2), + "full_seq_len": 4, + "full_vocab_size": 5, + }, + ] + } + ) + + out, fallback_count = rebuild_teacher_full_logits_from_ipc( + entries, cp_group=None, device="cpu" + ) + + assert fallback_count == 2 + assert torch.equal(out, torch.cat([vocab_left[0], vocab_right[0]], dim=-1)) + + def test_single_row_fallback_adds_batch_view_without_copy(self, monkeypatch): + assembled = torch.arange(4 * 5, dtype=torch.float32).reshape(4, 5) + entry = self._entry(0) + # TP-sharded metadata prevents the full-vocabulary zero-copy path. + entry["teacher_shards"][0]["vocab_end_index"] = 3 + monkeypatch.setattr( + "nemo_rl.algorithms.x_token.loss_utils.assemble_teacher_logits_from_shards", + lambda *_args, **_kwargs: assembled, + ) + + out, fallback_count = rebuild_teacher_full_logits_from_ipc( + [entry], cp_group=None, device=0 + ) + + assert fallback_count == 1 + assert out.shape == (1, 4, 5) + assert out.data_ptr() == assembled.data_ptr() + assert torch.equal(out[0], assembled) + + +class TestRebuildSparseTeacherLogits: + @staticmethod + def _shard(prefix: str, start: int, seq_len: int, full_seq_len: int) -> dict: + return { + "topk_logits_ipc": f"{prefix}-logits", + "topk_indices_ipc": f"{prefix}-indices", + "log_z_ipc": f"{prefix}-logz", + "gt_in_topk_ipc": f"{prefix}-gt", + "topk_shape": (seq_len, 2), + "global_seq_start": start, + "full_seq_len": full_seq_len, + } + + def test_orders_cp_shards_and_preserves_dtypes(self, monkeypatch): + tensors = { + "a-logits": torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + "a-indices": torch.tensor([[1, 3], [2, 4]], dtype=torch.int64), + "a-logz": torch.tensor([5.0, 6.0]), + "a-gt": torch.tensor([1, 0], dtype=torch.int32), + "b-logits": torch.tensor([[7.0, 8.0]]), + "b-indices": torch.tensor([[5, 6]], dtype=torch.int64), + "b-logz": torch.tensor([9.0]), + "b-gt": torch.tensor([0], dtype=torch.int32), + } + monkeypatch.setattr( + "nemo_rl.models.policy.utils.rebuild_cuda_tensor_from_ipc", + lambda handle, _device: tensors[handle], + ) + entries = [ + { + "teacher_shards": [ + self._shard("b", 2, 1, 3), + self._shard("a", 0, 2, 3), + ] + } + ] + + logits, indices, log_z, gt_in_topk = rebuild_teacher_sparse_logits_from_ipc( + entries, device=0 + ) + + assert logits.shape == (1, 3, 2) + assert logits.dtype == torch.float32 + assert torch.equal(logits[0, :, 0], torch.tensor([1.0, 3.0, 7.0])) + assert indices.dtype == torch.int32 + assert torch.equal(indices[0], torch.tensor([[1, 3], [2, 4], [5, 6]])) + assert log_z.dtype == torch.float32 + assert torch.equal(log_z, torch.tensor([[5.0, 6.0, 9.0]])) + assert gt_in_topk is not None and gt_in_topk.dtype == torch.bool + assert torch.equal(gt_in_topk, torch.tensor([[True, False, False]])) + + def test_rejects_noncontiguous_cp_coverage(self, monkeypatch): + monkeypatch.setattr( + "nemo_rl.models.policy.utils.rebuild_cuda_tensor_from_ipc", + lambda _handle, _device: torch.zeros(1), + ) + entries = [ + {"teacher_shards": [self._shard("gap", 1, 1, 2)]}, + ] + with pytest.raises(ValueError, match="expected start 0, got 1"): + rebuild_teacher_sparse_logits_from_ipc(entries, device=0) + class TestAssembleTeacherLogitsFromShards: """CPU coverage of the heterogeneous teacher-TP/CP -> student-CP reassembly. @@ -682,6 +1071,57 @@ def _build_tp2cp2(full): ) return shards, storage_map + @staticmethod + def _build_compact_tp2cp2(full, valid_seq_len): + """Split one valid-prefix row into flat teacher TP2xCP2 IPC slabs.""" + full_seq_len, full_vocab_size = full.shape + seq_mid, vocab_mid = full_seq_len // 2, full_vocab_size // 2 + storage_map = {} + shards = [] + for cp, (seq_start, seq_end) in enumerate( + [(0, seq_mid), (seq_mid, full_seq_len)] + ): + stored_seq_len = max(0, min(valid_seq_len, seq_end) - seq_start) + for tp, (vocab_start, vocab_end) in enumerate( + [(0, vocab_mid), (vocab_mid, full_vocab_size)] + ): + payload = f"compact-cp{cp}tp{tp}" + if stored_seq_len > 0: + storage_map[payload] = full[ + seq_start : seq_start + stored_seq_len, + vocab_start:vocab_end, + ].clone() + payload_ipc = payload + else: + payload_ipc = None + local_vocab_size = vocab_end - vocab_start + shards.append( + { + "payload_ipc": payload_ipc, + "buf_idx": 0, + "sample_index_in_buf": 0, + "ipc_layout": DENSE_TEACHER_IPC_FLAT_LAYOUT, + "storage_shape": (stored_seq_len, local_vocab_size), + "storage_token_offset": 0, + "storage_used_tokens": stored_seq_len, + "storage_capacity_tokens": stored_seq_len, + "stored_seq_len": stored_seq_len, + "valid_seq_len": valid_seq_len, + "vocab_start_index": vocab_start, + "vocab_end_index": vocab_end, + "global_seq_start": seq_start, + "actual_shape": (seq_end - seq_start, local_vocab_size), + "full_seq_len": full_seq_len, + "full_vocab_size": full_vocab_size, + "tp_rank": tp, + "tp_size": 2, + "cp_rank": cp, + "cp_size": 2, + "dtype": torch.float32, + } + ) + return shards, storage_map + @pytest.mark.parametrize("student_cp_rank", [0, 1]) def test_tp2cp2_reassembly(self, monkeypatch, student_cp_rank): full = torch.arange(8 * 6, dtype=torch.float32).reshape(8, 6) @@ -700,6 +1140,54 @@ def test_tp2cp2_reassembly(self, monkeypatch, student_cp_rank): assert out.shape == (4, 6) assert torch.equal(out, full[lo:hi, :]) + @pytest.mark.parametrize( + "student_cp_size,student_cp_rank", + [(1, 0), (2, 0), (2, 1), (4, 2), (4, 3)], + ) + def test_compact_tp2cp2_reassembly_zero_fills_across_cp_remap( + self, monkeypatch, student_cp_size, student_cp_rank + ): + full = torch.arange(8 * 6, dtype=torch.float32).reshape(8, 6) + valid_seq_len = 5 + shards, storage_map = self._build_compact_tp2cp2(full, valid_seq_len) + monkeypatch.setattr( + "nemo_rl.models.policy.utils.rebuild_cuda_tensor_from_ipc", + lambda payload, _device: storage_map[payload], + ) + + out = assemble_teacher_logits_from_shards( + shards, + student_cp_rank=student_cp_rank, + student_cp_size=student_cp_size, + device="cpu", + ) + + expected = full.clone() + expected[valid_seq_len:].zero_() + local_seq_len = full.shape[0] // student_cp_size + seq_start = student_cp_rank * local_seq_len + assert out.shape == (local_seq_len, full.shape[1]) + assert torch.equal(out, expected[seq_start : seq_start + local_seq_len]) + + def test_compact_zero_teacher_cp_shard_is_not_mapped(self, monkeypatch): + full = torch.arange(8 * 6, dtype=torch.float32).reshape(8, 6) + shards, storage_map = self._build_compact_tp2cp2(full, valid_seq_len=4) + + def rebuild(payload, _device): + assert payload is not None + return storage_map[payload] + + monkeypatch.setattr( + "nemo_rl.models.policy.utils.rebuild_cuda_tensor_from_ipc", rebuild + ) + out = assemble_teacher_logits_from_shards( + shards, student_cp_rank=0, student_cp_size=1, device="cpu" + ) + + expected = full.clone() + expected[4:].zero_() + assert torch.equal(out, expected) + def test_non_divisible_seq_len_asserts(self): # full_seq_len=7 not divisible by student_cp_size=2 -> guard fires before # the out-of-bounds dest write the assert protects against. @@ -864,6 +1352,346 @@ def run_cp(self): if torch.distributed.is_initialized(): torch.distributed.destroy_process_group() + def run_CSD(self): + """TP/CP parallelism-invariance for the CSD cross-tokenizer partition KL. + + Builds a full-tensor single-GPU reference (loss + student-logit grad) and + a sharded run — student vocab-sharded over TP, student/teacher sequence + contiguous-block-sharded over CP — of + ``_compute_csd_loss``, then asserts: + + * each rank's loss == ``reference_loss / cp_size``. The in-loss + ``global_valid_chunks`` normalizer is the DP×CP-reduced chunk count + (WORLD-reduce ÷ tp_world), so TP is *replicated* (no TP factor, + matching CE's ``global_valid_toks``) while each CP rank holds a + ``1/cp_size`` share that sums back to the full loss; and + * the student-logit gradient, reconstructed from the shards, equals + ``reference_grad`` exactly. The CP contiguous-gather backward + all-reduces the per-rank grad to full magnitude (the CP factor + cancels) and TP carries no loss-scale factor, so nothing is left over. + + The reference passes ``global_valid_chunks`` explicitly (so it does *not* + WORLD-reduce) with ``tp_group=cp_group=None`` (gathers are no-ops), so it + is a purely local per-rank computation with no collectives. + """ + import os + import tempfile + + from nemo_rl.algorithms.loss.csd import ( + CSDLossFn as CrossTokenizerDistillationLossFn, + ) + from nemo_rl.algorithms.x_token.loss_utils import LocalizedAlignment + + try: + torch.distributed.init_process_group(backend="nccl") + rank = int(os.environ["RANK"]) + ws = int(os.environ["WORLD_SIZE"]) + tp_size, cp_size = self.tp_size, self.cp_size + assert ws == tp_size * cp_size + device = torch.device("cuda") + # NamedSharding layout is ``arange(ws).reshape(tp, cp)`` so + # ``rank == tp_idx * cp_size + cp_idx``. + tp_idx, cp_idx = rank // cp_size, rank % cp_size + + # TP / CP subgroups. Every rank must invoke every ``new_group`` + # (it is a collective over the default group), keeping only its own. + tp_group = None + for c in range(cp_size): + g = torch.distributed.new_group( + ranks=[t * cp_size + c for t in range(tp_size)] + ) + if c == cp_idx: + tp_group = g + cp_group = None + for t in range(tp_size): + g = torch.distributed.new_group( + ranks=[t * cp_size + c for c in range(cp_size)] + ) + if t == tp_idx: + cp_group = g + if tp_size == 1: + tp_group = None + if cp_size == 1: + cp_group = None + + # ---- Fixture (identical on every rank; sharded below) ---------- + # Qualification-shaped CSD fixture with repeated common, 1-to-2, and + # 2-to-1 chunks. Counts cross the fixed common streaming cap (64 + # rows); shifted teacher predictors live on both CP owners. Common + # chunks use JSD and mismatch-last uses KL. + torch.manual_seed(0) + B, S, T, v_s, v_t = 2, 6, 6, 16, 16 + temp = 1.0 + student_ids = torch.tensor( + [[0, 5, 0, 9, 5, 0], [1, 6, 0, 9, 6, 0]], dtype=torch.long + ) + teacher_ids = torch.tensor( + [[0, 8, 3, 0, 5, 0], [1, 8, 4, 0, 6, 0]], dtype=torch.long + ) + chunks = [(0, 1, 0, 1)] * 33 + [(1, 2, 1, 3), (3, 5, 4, 5)] * 5 + max_pairs = len(chunks) + s_spans = torch.zeros(B, max_pairs, 2, dtype=torch.long) + t_spans = torch.zeros(B, max_pairs, 2, dtype=torch.long) + pair_valid = torch.zeros(B, max_pairs, dtype=torch.bool) + num_chunks = torch.full((B,), max_pairs, dtype=torch.long) + for b in range(B): + for k, (ss, se, ts, te) in enumerate(chunks): + s_spans[b, k, 0], s_spans[b, k, 1] = ss, se + t_spans[b, k, 0], t_spans[b, k, 1] = ts, te + pair_valid[b, k] = True + student_full = torch.randn(B, S, v_s) # CPU -> deterministic + # Huge padded columns prove dense CSD slices to the real tokenizer + # vocabulary before exact logZ and CP-owner lookups. + teacher_full = torch.randn(B, T, v_t + 2) + teacher_full[..., v_t:] = 1.0e4 + gvc = float(B * len(chunks)) # total valid chunks (single-rank) + + student_ids = student_ids.to(device) + teacher_ids = teacher_ids.to(device) + s_spans, t_spans = s_spans.to(device), t_spans.to(device) + pair_valid, num_chunks = pair_valid.to(device), num_chunks.to(device) + student_full, teacher_full = ( + student_full.to(device), + teacher_full.to(device), + ) + + # Tiny common + bidirectional prefix-support tables. + tmpdir = tempfile.mkdtemp() + subtoks = torch.full((v_s, 3), -1, dtype=torch.long) + lengths = torch.zeros((v_s,), dtype=torch.long) + for s in range(10): + subtoks[s, 0], lengths[s] = s, 1 + subtoks[10, 0], subtoks[10, 1], lengths[10] = 8, 3, 2 + subtoks[11, 0], subtoks[11, 1], lengths[11] = 8, 4, 2 + for s in range(12, v_s): + subtoks[s, 0], subtoks[s, 1], lengths[s] = s % 10, (s + 3) % 10, 2 + fwd_path = os.path.join(tmpdir, "fwd.pt") + rev_path = os.path.join(tmpdir, "rev.pt") + torch.save({"subtoks": subtoks, "lengths": lengths}, fwd_path) + reverse_subtoks = torch.full((v_t, 3), -1, dtype=torch.long) + reverse_lengths = torch.zeros((v_t,), dtype=torch.long) + for t in range(10): + reverse_subtoks[t, 0], reverse_lengths[t] = t, 1 + reverse_subtoks[10, 0], reverse_subtoks[10, 1], reverse_lengths[10] = ( + 9, + 3, + 2, + ) + reverse_subtoks[11, 0], reverse_subtoks[11, 1], reverse_lengths[11] = ( + 9, + 4, + 2, + ) + for t in range(12, v_t): + reverse_subtoks[t, 0], reverse_subtoks[t, 1], reverse_lengths[t] = ( + t % 10, + (t + 3) % 10, + 2, + ) + torch.save( + {"subtoks": reverse_subtoks, "lengths": reverse_lengths}, rev_path + ) + + cfg = { + "temperature": temp, + "vocab_topk": 8, + "reverse_kl": False, + "kl_loss_weight": 1.0, + "ce_loss_scale": 1.0, + "dynamic_loss_scaling": False, + "kd_loss_mode": "sum", + "normalize_teacher_by_vocab": False, + "alpha": 1.0, + "sum_weights_metric": None, + "student_vocab_size": v_s, + "teacher_vocab_sizes": [v_t], + "projection_matrix_paths": ["dummy.pt"], + "teacher_weights": [1.0], + "common_indices_from_subtoks": True, + "student_to_teacher_table_paths": [fwd_path], + "teacher_to_student_table_paths": [rev_path], + "kl_chunk_shift": True, + "csd_loss_fn": "jsd", + "csd_last_pos_loss_fn": "kl", + "csd_jsd_beta": 0.5, + "csd_noise_filter_topk": 0, + "teacher_topk_ipc_k": 0, + "teacher_topk_ipc_support_mode": "row_topk", + "teacher_topk_ipc_keep_realized": True, + } + loss_fn = CrossTokenizerDistillationLossFn(cfg) + + def _align(student_input_ids, teacher_input_ids): + return LocalizedAlignment( + sample_mask=torch.ones(B, dtype=torch.bool, device=device), + pair_valid=pair_valid, + student_input_ids=student_input_ids, + teacher_input_ids=teacher_input_ids, + student_spans=s_spans, + teacher_spans=t_spans, + num_chunks=num_chunks, + ) + + # ---- Single-GPU reference (explicit gvc => no WORLD-reduce; tp/cp + # None => gathers are no-ops => purely local, no collectives). -- + ref_logits = student_full.clone().requires_grad_(True) + ref_loss, ref_metrics = loss_fn._compute_csd_loss( + 0, + ref_logits, + teacher_full[..., :v_t].clone(), + _align(student_ids, teacher_ids), + teacher_vocab_size=v_t, + tp_group=None, + cp_group=None, + global_valid_chunks=torch.tensor(gvc, device=device), + ) + ref_loss.backward() + ref_grad = ref_logits.grad.detach().clone() # [B, S, v_s] + + # ---- Sharded run: CP contiguous seq block x TP vocab slice. ------- + sblk, tblk, vblk = S // cp_size, T // cp_size, v_s // tp_size + s_lo, t_lo, v_lo = cp_idx * sblk, cp_idx * tblk, tp_idx * vblk + student_shard = ( + student_full[:, s_lo : s_lo + sblk, v_lo : v_lo + vblk] + .detach() + .clone() + .requires_grad_(True) + ) + # Teacher is full-vocab (only CP-sharded on the sequence). + teacher_shard = teacher_full[:, t_lo : t_lo + tblk, :].detach().clone() + import nemo_rl.algorithms.loss.csd as csd_loss_module + + original_cp_gather = csd_loss_module.allgather_cp_contiguous_tensor + + def reject_dense_teacher_sequence_gather(tensor, group, seq_dim=1): + if ( + group is cp_group + and tensor.ndim == 3 + and not tensor.requires_grad + and tensor.shape[0] == B + and tensor.shape[1] == tblk + and tensor.shape[-1] >= v_t + ): + raise AssertionError( + "dense teacher [B,T/CP,V] entered a full CP sequence gather" + ) + return original_cp_gather(tensor, group, seq_dim) + + csd_loss_module.allgather_cp_contiguous_tensor = ( + reject_dense_teacher_sequence_gather + ) + try: + sh_loss, sh_metrics = loss_fn._compute_csd_loss( + 0, + student_shard, + teacher_shard, + _align( + student_ids[:, s_lo : s_lo + sblk].contiguous(), + teacher_ids[:, t_lo : t_lo + tblk].contiguous(), + ), + teacher_vocab_size=v_t, + tp_group=tp_group, + cp_group=cp_group, + global_valid_chunks=None, # => in-loss WORLD-reduce over ws + ) + finally: + csd_loss_module.allgather_cp_contiguous_tensor = original_cp_gather + sh_loss.backward() + sh_grad = student_shard.grad.detach().clone() # [B, S/cp, v_s/tp] + + # Loss invariance: every rank == reference / cp_size (TP replicated; + # each CP rank holds a 1/cp_size share of the full loss). + torch.testing.assert_close( + sh_loss.detach(), + ref_loss.detach() / cp_size, + rtol=1e-4, + atol=1e-4, + ) + + # Grad invariance: reconstruct the full student grad from the shards; + # it equals the single-GPU reference exactly (CP contributions + # all-reduce to full magnitude, TP carries no loss-scale factor). + gathered = [torch.empty_like(sh_grad) for _ in range(ws)] + torch.distributed.all_gather(gathered, sh_grad.contiguous()) + recon = torch.zeros_like(ref_grad) + for r in range(ws): + r_tp, r_cp = r // cp_size, r % cp_size + recon[ + :, r_cp * sblk : (r_cp + 1) * sblk, r_tp * vblk : (r_tp + 1) * vblk + ] = gathered[r] + torch.testing.assert_close(recon, ref_grad, rtol=1e-4, atol=1e-4) + + # Production supplies the full-batch chunk denominator. The raw + # replicated CP loss then equals the CP1 reference; MCore's CP loss + # wrapper contributes the 1/CP factor before backward. + explicit_student_shard = ( + student_full[:, s_lo : s_lo + sblk, v_lo : v_lo + vblk] + .detach() + .clone() + .requires_grad_(True) + ) + explicit_loss, explicit_metrics = loss_fn._compute_csd_loss( + 0, + explicit_student_shard, + teacher_shard, + _align( + student_ids[:, s_lo : s_lo + sblk].contiguous(), + teacher_ids[:, t_lo : t_lo + tblk].contiguous(), + ), + teacher_vocab_size=v_t, + tp_group=tp_group, + cp_group=cp_group, + global_valid_chunks=torch.tensor(gvc, device=device), + ) + torch.testing.assert_close( + explicit_loss.detach(), ref_loss.detach(), rtol=1e-4, atol=1e-4 + ) + (explicit_loss / float(cp_size)).backward() + explicit_grad = explicit_student_shard.grad.detach().clone() + gathered_explicit = [torch.empty_like(explicit_grad) for _ in range(ws)] + torch.distributed.all_gather( + gathered_explicit, + explicit_grad.contiguous(), + ) + explicit_recon = torch.zeros_like(ref_grad) + for r in range(ws): + r_tp, r_cp = r // cp_size, r % cp_size + explicit_recon[ + :, + r_cp * sblk : (r_cp + 1) * sblk, + r_tp * vblk : (r_tp + 1) * vblk, + ] = gathered_explicit[r] + torch.testing.assert_close( + explicit_recon, + ref_grad, + rtol=1e-4, + atol=1e-4, + ) + sh_metrics = explicit_metrics + assert sh_metrics["num_common_chunks"] > 0 + assert sh_metrics["num_mismatch_chunks"] > 0 + for metric_name in ( + "kl_common_per_chunk", + "kl_partition_last_per_chunk", + "kl_mismatch_scaled_per_chunk", + "top1_acc_per_chunk", + ): + assert sh_metrics[metric_name] == pytest.approx( + ref_metrics[metric_name], rel=1.0e-4, abs=1.0e-4 + ) + for metric_name in ( + "num_common_chunks", + "num_mismatch_chunks", + "num_valid_samples", + ): + assert sh_metrics[metric_name] == ref_metrics[metric_name] + return {"success": True, "error": None} + except Exception: + return {"success": False, "error": traceback.format_exc()} + finally: + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + _ACTOR_FQN = f"{XtokenShardTestActor.__module__}.XtokenShardTestActor" @@ -1673,3 +2501,18 @@ def test_same_tokenizer_kd_cp2_gloo_matches_single_process(tmp_path: Path): ) assert sorted(results) == list(range(world_size)) assert [proc.exitcode for proc in procs] == [0] * world_size + + +def test_csd_tp2cp1(register_actor): + """CSD partition-KL is invariant under TP=2 vocab sharding.""" + _run(register_actor, tp_size=2, cp_size=1, method="run_CSD") + + +def test_csd_tp1cp2(register_actor): + """CSD partition-KL is invariant under CP=2 sequence sharding.""" + _run(register_actor, tp_size=1, cp_size=2, method="run_CSD") + + +def test_csd_tp2cp2(register_actor): + """CSD partition-KL is invariant under composed TP=2 x CP=2 sharding.""" + _run(register_actor, tp_size=2, cp_size=2, method="run_CSD") diff --git a/tests/unit/algorithms/x_token/test_teacher_row_readers.py b/tests/unit/algorithms/x_token/test_teacher_row_readers.py new file mode 100644 index 00000000000..0a68dbb4646 --- /dev/null +++ b/tests/unit/algorithms/x_token/test_teacher_row_readers.py @@ -0,0 +1,373 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU coverage for the sparse and dense teacher row readers.""" + +import copy + +import pytest +import torch + +from nemo_rl.algorithms.x_token.loss_utils import ( + DenseTeacherRowReader, + SparseTeacherRowReader, + gather_sparse_teacher_rows_from_ipc, +) +from nemo_rl.distributed.model_utils import ( + _get_tokens_on_this_cp_rank, + cp_load_balanced_segments, + distributed_vocab_topk_logz_force, +) +from nemo_rl.models.policy.utils import DENSE_TEACHER_IPC_FLAT_LAYOUT + +CPU = torch.device("cpu") +BATCH, SEQ, VOCAB, K, TEMPERATURE = 2, 8, 12, 3, 0.7 + + +def _canonical_order(row_logits: torch.Tensor) -> list[int]: + """Token ids by descending logit, ties broken by the lowest id.""" + return sorted(range(row_logits.numel()), key=lambda v: (-float(row_logits[v]), v)) + + +def _reference(): + torch.manual_seed(7) + logits = torch.randn(BATCH, SEQ, VOCAB) + # Row (1, 5) has an exact cutoff tie (ids 3 and 9 at rank K), so the + # evicted slot is ambiguous by value and must follow the recorded tail id. + logits[1, 5, :] = torch.linspace(-1.0, 1.0, VOCAB) + logits[1, 5, 3] = logits[1, 5, -3] + force = torch.full((BATCH, SEQ, 2), -1) + for b in range(BATCH): + for t in range(SEQ): + order = _canonical_order(logits[b, t]) + force[b, t, 0] = order[0] # a natural top-k label + force[b, t, 1] = order[-1] # a label outside the natural top-k + return logits, force + + +def _sparse_payload(cp_size: int = 2): + """Export the reference with the real primitive and split it like Megatron.""" + logits, force = _reference() + fields = distributed_vocab_topk_logz_force( + logits, + force, + K, + None, + vocab_start_index=0, + vocab_end_index=VOCAB, + real_vocab_size=VOCAB, + temperature=TEMPERATURE, + ) + names = ( + "topk_logits", + "topk_indices", + "log_z", + "natural_tail_indices", + "forced_logits", + "forced_indices", + "forced_in_topk", + ) + full = dict(zip(names, fields)) + full["log_z"] = full["log_z"].unsqueeze(-1) + entries = [{"teacher_shards": []} for _ in range(BATCH)] + for cp_rank in range(cp_size): + local = { + name: _get_tokens_on_this_cp_rank(value, cp_rank, cp_size, seq_dim=1) + .unsqueeze(0) + .contiguous() + for name, value in full.items() + } + local_seq = SEQ // cp_size + for b in range(BATCH): + entries[b]["teacher_shards"].append( + { + "transport": "sparse_topk", + **{f"{name}_ipc": tensor for name, tensor in local.items()}, + "buf_idx": 0, + "sample_index_in_buf": b, + "topk_shape": (local_seq, K), + "forced_shape": (local_seq, 2), + "scalar_shape": (local_seq,), + "k": K, + "force_width": 2, + "temperature": TEMPERATURE, + "real_vocab_size": VOCAB, + "forced_in_topk_k": K, + "cp_rank": cp_rank, + "cp_size": cp_size, + "full_seq_len": SEQ, + "sequence_segments": cp_load_balanced_segments( + SEQ, cp_rank, cp_size + ), + } + ) + return entries, logits, force + + +def _expected_support(row_logits: torch.Tensor, label: int): + natural = _canonical_order(row_logits)[:K] + if label not in natural: + natural = natural[: K - 1] + [label] + ids = sorted(natural) + return ids, row_logits[ids] + + +@pytest.mark.parametrize("cp_size", [1, 2, 4]) +def test_sparse_reader_materializes_natural_or_forced_support(cp_size): + entries, logits, force = _sparse_payload(cp_size) + reader = SparseTeacherRowReader(entries, device=CPU) + requests = [ + (b, t, slot) for b in range(BATCH) for t in range(SEQ) for slot in (0, 1) + ] + batch = torch.tensor([b for b, _, _ in requests]) + positions = torch.tensor([t for _, t, _ in requests]) + labels = torch.stack([force[b, t, slot] for b, t, slot in requests]) + + values, ids, log_z, in_topk = reader.gather_rows(batch, positions, labels) + + assert (reader.k, reader.full_seq_len, reader.force_width) == (K, SEQ, 2) + assert values.dtype == torch.float32 and ids.dtype == torch.int32 + for row, (b, t, slot) in enumerate(requests): + expected_ids, expected_values = _expected_support( + logits[b, t], int(force[b, t, slot]) + ) + assert ids[row].tolist() == expected_ids + torch.testing.assert_close(values[row], expected_values) + assert bool(in_topk[row]) == (slot == 0) + torch.testing.assert_close( + log_z, torch.logsumexp(logits[batch, positions] / TEMPERATURE, dim=-1) + ) + + +def test_sparse_reader_reuses_routes_and_backings_across_bounded_gathers(monkeypatch): + # A tensor-backed CPU payload must not consult CUDA. + monkeypatch.setattr( + torch.cuda, + "current_device", + lambda: (_ for _ in ()).throw(AssertionError("CUDA was consulted")), + ) + entries, _, force = _sparse_payload() + reader = SparseTeacherRowReader(entries, device=CPU) + + reader.gather_rows( + torch.tensor([1, 0]), torch.tensor([5, 7]), force[[1, 0], [5, 7], 0] + ) + cached = {key: tensor.data_ptr() for key, tensor in reader._backing_tensors.items()} + reader.gather_rows( + torch.tensor([1, 0]), torch.tensor([3, 6]), force[[1, 0], [3, 6], 1] + ) + + # Seven fields per CP rank's slab, shared by every sample of that rank. + assert len(cached) == 7 * 2 + assert cached == { + key: tensor.data_ptr() for key, tensor in reader._backing_tensors.items() + } + reader.close() + assert reader._backing_tensors == {} + + +def test_one_shot_helper_matches_reader(): + entries, _, force = _sparse_payload() + batch = torch.tensor([1, 0, 1, 0]) + positions = torch.tensor([5, 7, 0, 3]) + labels = force[batch, positions, 1] + expected = SparseTeacherRowReader(entries, device=CPU).gather_rows( + batch, positions, labels + ) + actual = gather_sparse_teacher_rows_from_ipc( + entries, batch, positions, labels, device=CPU + ) + for actual_tensor, expected_tensor in zip(actual, expected): + torch.testing.assert_close(actual_tensor, expected_tensor) + + +@pytest.mark.parametrize("failure", ["overlap", "gap", "local_gap"]) +def test_sparse_reader_rejects_invalid_segment_coverage(failure): + entries, _, _ = _sparse_payload() + broken = copy.deepcopy(entries) + shards = broken[0]["teacher_shards"] + if failure == "overlap": + shards[1]["sequence_segments"] = [(0, 1, 2), (2, 4, 2)] + elif failure == "gap": + shards[1]["sequence_segments"] = [(0, 3, 1), (1, 4, 3)] + else: + shards[0]["sequence_segments"] = [(0, 0, 2), (3, 6, 1)] + + with pytest.raises(ValueError, match="cover"): + SparseTeacherRowReader(broken, device=CPU) + + +def test_sparse_reader_rejects_unknown_transport_and_unforced_labels(): + entries, logits, force = _sparse_payload() + legacy = copy.deepcopy(entries) + legacy[0]["teacher_shards"][0]["transport"] = "dense" + with pytest.raises(ValueError, match="Unsupported sparse teacher transport"): + SparseTeacherRowReader(legacy, device=CPU) + + reader = SparseTeacherRowReader(entries, device=CPU) + unforced = int(next(v for v in range(VOCAB) if v not in force[0, 2].tolist())) + with pytest.raises(ValueError, match="exactly once in the forced sidecar"): + reader.gather_rows( + torch.tensor([0]), torch.tensor([2]), torch.tensor([unforced]) + ) + with pytest.raises(ValueError, match="real-vocabulary ids"): + reader.gather_rows(torch.tensor([0]), torch.tensor([2]), torch.tensor([VOCAB])) + with pytest.raises(ValueError, match="exactly once"): + reader.gather_rows(torch.tensor([0]), torch.tensor([SEQ]), force[0, 0, :1]) + + +def test_sparse_reader_rejects_inconsistent_forced_copy(): + entries, _, force = _sparse_payload(cp_size=1) + shard = entries[0]["teacher_shards"][0] + # Row 0's slot-0 label is a natural id; corrupt its forced logit. + shard["forced_logits_ipc"] = shard["forced_logits_ipc"].clone() + shard["forced_logits_ipc"][0, 0, 0, 0] += 1.0 + reader = SparseTeacherRowReader(entries, device=CPU) + with pytest.raises(ValueError, match="different raw logits"): + reader.gather_rows(torch.tensor([0]), torch.tensor([0]), force[0, 0, :1]) + + +def _dense_payload(*, tp_size: int, cp_size: int, compact_valid_len=None): + """Dense teacher records in the Megatron/Automodel sender formats.""" + torch.manual_seed(3) + logits = torch.randn(BATCH, SEQ, VOCAB) + local_seq, local_vocab = SEQ // cp_size, VOCAB // tp_size + entries = [{"teacher_shards": []} for _ in range(BATCH)] + storages = [] + for cp_rank in range(cp_size): + for tp_rank in range(tp_size): + seq = slice(cp_rank * local_seq, (cp_rank + 1) * local_seq) + vocab = slice(tp_rank * local_vocab, (tp_rank + 1) * local_vocab) + common = { + "tp_rank": tp_rank, + "tp_size": tp_size, + "cp_rank": cp_rank, + "cp_size": cp_size, + "vocab_start_index": vocab.start, + "vocab_end_index": vocab.stop, + "global_seq_start": seq.start, + "full_seq_len": SEQ, + "full_vocab_size": VOCAB, + "actual_shape": (local_seq, local_vocab), + } + if compact_valid_len is None: + storage = logits[:, seq, vocab].unsqueeze(0).contiguous() + for b in range(BATCH): + entries[b]["teacher_shards"].append( + { + **common, + "payload_ipc": ("rect", cp_rank, tp_rank), + "buf_idx": 0, + "sample_index_in_buf": b, + } + ) + else: + stored = [ + max(0, min(valid, seq.stop) - seq.start) + for valid in compact_valid_len + ] + storage = torch.cat( + [ + logits[b, seq.start : seq.start + stored[b], vocab] + for b in range(BATCH) + ] + ).contiguous() + offset = 0 + for b in range(BATCH): + entries[b]["teacher_shards"].append( + { + **common, + "payload_ipc": ("compact", cp_rank, tp_rank), + "buf_idx": 0, + "sample_index_in_buf": b, + "ipc_layout": DENSE_TEACHER_IPC_FLAT_LAYOUT, + "storage_shape": tuple(storage.shape), + "storage_token_offset": offset, + "storage_used_tokens": int(storage.shape[0]), + "storage_capacity_tokens": int(storage.shape[0]), + "stored_seq_len": stored[b], + "valid_seq_len": compact_valid_len[b], + "dtype": storage.dtype, + } + ) + offset += stored[b] + storages.append(((cp_rank, tp_rank), storage)) + return entries, logits, dict(storages) + + +def _patch_rebuild(monkeypatch, storages): + opened = [] + + def rebuild(handle, _device): + opened.append(handle) + return storages[(handle[1], handle[2])] + + monkeypatch.setattr( + "nemo_rl.models.policy.utils.rebuild_cuda_tensor_from_ipc", rebuild + ) + return opened + + +@pytest.mark.parametrize("tp_size,cp_size", [(1, 1), (2, 1), (2, 2), (3, 4)]) +def test_dense_reader_stitches_requested_rows_across_tp_cp( + monkeypatch, tp_size, cp_size +): + entries, logits, storages = _dense_payload(tp_size=tp_size, cp_size=cp_size) + opened = _patch_rebuild(monkeypatch, storages) + reader = DenseTeacherRowReader(entries, device=torch.device("cuda", 0)) + reader.device = CPU # CPU stand-in for the consumer GPU + batch = torch.tensor([1, 0, 1, 1, 0]) + positions = torch.tensor([7, 0, 3, 3, 5]) + + rows = reader.gather_rows(batch, positions) + torch.testing.assert_close(rows, logits[batch, positions]) + assert (reader.full_seq_len, reader.full_vocab_size) == (SEQ, VOCAB) + + reader.gather_rows(torch.tensor([0, 1]), torch.tensor([6, 2])) + # Each producer storage is opened at most once per reader. + assert len(opened) == len(set(opened)) <= tp_size * cp_size + reader.close() + assert reader._storage == {} + + +def test_dense_reader_reads_compact_suffix_as_zero(monkeypatch): + entries, logits, storages = _dense_payload( + tp_size=2, cp_size=2, compact_valid_len=[5, 2] + ) + _patch_rebuild(monkeypatch, storages) + reader = DenseTeacherRowReader(entries, device=torch.device("cuda", 0)) + reader.device = CPU + batch = torch.tensor([0, 0, 0, 1, 1]) + positions = torch.tensor([4, 5, 7, 1, 6]) + + rows = reader.gather_rows(batch, positions) + + torch.testing.assert_close(rows[0], logits[0, 4]) + torch.testing.assert_close(rows[3], logits[1, 1]) + for zero_row in (1, 2, 4): + assert torch.equal(rows[zero_row], torch.zeros(VOCAB)) + + +def test_dense_reader_validates_grid_and_positions(monkeypatch): + entries, _, storages = _dense_payload(tp_size=2, cp_size=2) + _patch_rebuild(monkeypatch, storages) + incomplete = copy.deepcopy(entries) + incomplete[0]["teacher_shards"].pop() + with pytest.raises(ValueError, match="incomplete TP/CP shard coverage"): + DenseTeacherRowReader(incomplete, device=CPU) + + reader = DenseTeacherRowReader(entries, device=CPU) + with pytest.raises(IndexError, match="position"): + reader.gather_rows(torch.tensor([0]), torch.tensor([SEQ])) + with pytest.raises(IndexError, match="batch index"): + reader.gather_rows(torch.tensor([BATCH]), torch.tensor([0])) diff --git a/tests/unit/algorithms/x_token/test_token_aligner.py b/tests/unit/algorithms/x_token/test_token_aligner.py index 7cd9eeee492..e956393b4af 100644 --- a/tests/unit/algorithms/x_token/test_token_aligner.py +++ b/tests/unit/algorithms/x_token/test_token_aligner.py @@ -288,6 +288,43 @@ def test_align_padding_positions_receive_sentinel_chunk_id(): assert bool((t_sizes[chunk_mask] > 0).all()) +def test_align_pair_valid_drops_padding_pairs(): + """In a right-padded batch with rows of different lengths, the pairs over + padding are invalid, while real two-sided pairs and the student-only + start-of-text pair (one-sided by alignment) stay valid. + """ + vocab = {0: "", 7: "", 1: "Hello", 2: "Ġworld", 3: "!"} + aligner = _make_aligner(vocab, vocab, bos_token_id=7, pad_token_id=0) + + # Only the student starts with ; row 1 is two tokens shorter. + student_ids = torch.tensor([[7, 1, 2, 3], [7, 1, 0, 0]], dtype=torch.long) + teacher_ids = torch.tensor([[1, 2, 3], [1, 0, 0]], dtype=torch.long) + student_offsets = torch.tensor( + [[[0, 0], [0, 5], [5, 11], [11, 12]], [[0, 0], [0, 5], [0, 0], [0, 0]]], + dtype=torch.long, + ) + teacher_offsets = torch.tensor( + [[[0, 5], [5, 11], [11, 12]], [[0, 5], [0, 0], [0, 0]]], dtype=torch.long + ) + + batch = aligner.align( + student_ids, + teacher_ids, + student_offsets=student_offsets, + teacher_offsets=teacher_offsets, + student_attention_mask=student_ids.ne(0), + teacher_attention_mask=teacher_ids.ne(0), + ) + + # Pair 0 is alone; row 1's pairs 2-3 pair padding with padding. + assert batch.student_chunk_id.tolist() == [[0, 1, 2, 3], [0, 1, -1, -1]] + assert batch.teacher_chunk_id.tolist() == [[1, 2, 3], [1, -1, -1]] + assert batch.pair_valid.tolist() == [ + [True, True, True, True], + [True, True, False, False], + ] + + def test_module_exposes_alignment_kernel(): """The offset kernel is reachable from the package namespace (callers use it directly instead of poking at name-mangled internals). diff --git a/tests/unit/algorithms/x_token/test_utils.py b/tests/unit/algorithms/x_token/test_utils.py index 04b6cd38f97..91cc2460607 100644 --- a/tests/unit/algorithms/x_token/test_utils.py +++ b/tests/unit/algorithms/x_token/test_utils.py @@ -67,6 +67,8 @@ def _check(self, **kw): teacher_cp=2, student_dp=1, teacher_dp=1, + student_pp=1, + teacher_pp=1, ) args.update(kw) assert_xtoken_ipc_node_local(**args) @@ -98,6 +100,19 @@ def test_multinode_group_not_node_aligned_raises(self): gpus_per_node=6, student_tp=2, student_cp=2, teacher_tp=2, teacher_cp=2 ) # 6 % 4 != 0 + def test_pp_counts_toward_the_model_parallel_group(self): + # tp*cp*pp = 8 > gpus_per_node = 4: the producing last stage would land + # on another node from the student ranks importing its IPC buffers. + with pytest.raises(AssertionError): + self._check(student_pp=2, teacher_pp=2) + + def test_pp_mismatch_between_teacher_and_student_raises(self): + with pytest.raises(AssertionError): + self._check(gpus_per_node=8, student_pp=2, teacher_pp=1) + + def test_multinode_matched_grid_with_pp_ok(self): + self._check(gpus_per_node=8, student_pp=2, teacher_pp=2) + class TestPadDistillationValBatch: def test_pads_tensors_lists_and_zeros_mask(self): diff --git a/tests/unit/algorithms/x_token/test_xtoken_off_policy_distillation.py b/tests/unit/algorithms/x_token/test_xtoken_off_policy_distillation.py index 2a722d72030..03c47a634fa 100644 --- a/tests/unit/algorithms/x_token/test_xtoken_off_policy_distillation.py +++ b/tests/unit/algorithms/x_token/test_xtoken_off_policy_distillation.py @@ -455,6 +455,8 @@ def test_teacher_aligner_config_explicit_values_serialize_and_stay_out_of_policy assert dumped["aligner"] == { "projection_matrix_path": "/tmp/projection.pt", "drop_first_assistant_chunk_kl": True, + "student_to_teacher_table_path": None, + "teacher_to_student_table_path": None, } assert "projection_matrix_path" not in dumped assert "aligner" not in teacher.policy_config() diff --git a/tests/unit/data/datasets/test_arrow_text_dataset.py b/tests/unit/data/datasets/test_arrow_text_dataset.py index f9b1b6bf4e1..d7be57a1d4d 100644 --- a/tests/unit/data/datasets/test_arrow_text_dataset.py +++ b/tests/unit/data/datasets/test_arrow_text_dataset.py @@ -173,6 +173,7 @@ def test_emits_when_threshold_met(self): assert packs[0]["messages"] == [ {"role": "assistant", "content": "aaaa\nbbbb\ncccc"} ] + assert packs[0]["sample_id"] == "legacy-pack:0|1|2" def test_pack_may_exceed_threshold(self): # Two rows of 100 chars each, threshold 10 → first row alone @@ -254,3 +255,80 @@ def test_task_name_forwarded(self): rows = [{"text": "hello world"}] packs = _packs_from(rows, chars=5, task_name="my-task") assert packs[0]["task_name"] == "my-task" + + +def test_sample_id_is_created_before_filtering(mixed_arrow_file): + from nemo_rl.data.datasets.response_datasets.arrow_text_dataset import ( + ArrowTextDataset, + ) + from nemo_rl.data.datasets.response_datasets.response_dataset import ( + resolve_source_sample_id, + ) + + ds = ArrowTextDataset( + data_files=mixed_arrow_file, + text_key="text", + characters_per_sample=None, + ) + sample_ids = [row["sample_id"] for row in ds.dataset] + # Source ordinals 2 and 4 were invalid and filtered. The surviving IDs + # retain their original ordinals rather than being renumbered 0,1,2. + assert sample_ids == [ + resolve_source_sample_id( + {}, + ordinal, + data_path=mixed_arrow_file, + subset=None, + split="train", + ) + for ordinal in (0, 1, 3) + ] + + +def test_arrow_dataset_preserves_upstream_uuid(tmp_path: Path): + from nemo_rl.data.datasets.response_datasets.arrow_text_dataset import ( + ArrowTextDataset, + ) + + data_path = tmp_path / "uuid.arrow" + _write_arrow_file( + data_path, + [ + { + "text": "hello", + "uuid": "550e8400-e29b-41d4-a716-446655440000", + } + ], + ) + + dataset = ArrowTextDataset(str(data_path)) + + assert dataset.dataset[0]["sample_id"] == ("550e8400-e29b-41d4-a716-446655440000") + + +def test_arrow_namespace_distinguishes_path_subset_and_split(monkeypatch): + from nemo_rl.data.datasets.response_datasets import arrow_text_dataset as module + + source = Dataset.from_list([{"text": "hello"}]) + monkeypatch.setattr( + module, + "load_dataset_from_path", + lambda data_path, subset, split: source, + ) + namespaces = [ + ("org/corpus", "subset-a", "train"), + ("org/corpus", "subset-b", "train"), + ("org/corpus", "subset-a", "validation"), + ("org/other-corpus", "subset-a", "train"), + ] + + sample_ids = { + module.ArrowTextDataset( + data_files=data_path, + subset=subset, + split=split, + ).dataset[0]["sample_id"] + for data_path, subset, split in namespaces + } + + assert len(sample_ids) == len(namespaces) diff --git a/tests/unit/data/datasets/test_oai_format_dataset.py b/tests/unit/data/datasets/test_oai_format_dataset.py index ef7b000c59b..3fb332e9faf 100644 --- a/tests/unit/data/datasets/test_oai_format_dataset.py +++ b/tests/unit/data/datasets/test_oai_format_dataset.py @@ -21,6 +21,9 @@ from nemo_rl.data.chat_templates import COMMON_CHAT_TEMPLATES from nemo_rl.data.datasets import load_response_dataset from nemo_rl.data.datasets.response_datasets import OpenAIFormatDataset +from nemo_rl.data.datasets.response_datasets.response_dataset import ( + resolve_source_sample_id, +) @pytest.fixture @@ -115,3 +118,48 @@ def test_message_formatting(sample_data, tokenizer): assert combined_message == "".join( message["content"] for message in first_example["messages"] ) + + +@pytest.mark.parametrize("use_preserving_dataset", [False, True]) +def test_identity_is_attached_before_both_openai_formatting_paths( + tmp_path, use_preserving_dataset +): + data_path = tmp_path / "chat.jsonl" + records = [ + { + "messages": [{"role": "assistant", "content": "sample-id"}], + "sample_id": "native-sample-id", + }, + { + "messages": [{"role": "assistant", "content": "id"}], + "id": "native-id", + }, + { + "messages": [{"role": "assistant", "content": "uuid"}], + "uuid": "550e8400-e29b-41d4-a716-446655440000", + }, + {"messages": [{"role": "assistant", "content": "synthesized"}]}, + ] + data_path.write_text( + "\n".join(json.dumps(record) for record in records), encoding="utf-8" + ) + + dataset = OpenAIFormatDataset( + str(data_path), + use_preserving_dataset=use_preserving_dataset, + subset="conversation-v2", + split="validation", + ) + + assert [row["sample_id"] for row in dataset.dataset] == [ + "native-sample-id", + "native-id", + "550e8400-e29b-41d4-a716-446655440000", + resolve_source_sample_id( + records[3], + 3, + data_path=str(data_path), + subset="conversation-v2", + split="validation", + ), + ] diff --git a/tests/unit/data/datasets/test_response_dataset.py b/tests/unit/data/datasets/test_response_dataset.py index cb3130709bb..3cbc6c92645 100644 --- a/tests/unit/data/datasets/test_response_dataset.py +++ b/tests/unit/data/datasets/test_response_dataset.py @@ -14,6 +14,7 @@ import json import tempfile +from pathlib import Path import pytest from datasets import Dataset @@ -29,6 +30,10 @@ _format_options, ) from nemo_rl.data.datasets.response_datasets.nemogym_dataset import NemoGymDataset +from nemo_rl.data.datasets.response_datasets.response_dataset import ( + ResponseDataset, + resolve_source_sample_id, +) from nemo_rl.data.interfaces import NemoGymSourceIdentity from nemo_rl.data.processors import PROCESSOR_REGISTRY @@ -147,10 +152,11 @@ def test_response_dataset(input_key, output_key, is_save_to_disk, file_ext, toke # check the first example first_example = dataset.dataset[0] - # only contains messages and task_name - assert len(first_example.keys()) == 2 + # Formatted response rows retain their durable source identity. + assert len(first_example.keys()) == 3 assert "messages" in first_example assert "task_name" in first_example + assert "sample_id" in first_example # check the combined message chat_template = "{% for message in messages %}{%- if message['role'] == 'system' %}{{'Context: ' + message['content'].strip()}}{%- elif message['role'] == 'user' %}{{' Question: ' + message['content'].strip() + ' Answer:'}}{%- elif message['role'] == 'assistant' %}{{' ' + message['content'].strip()}}{%- endif %}{% endfor %}" @@ -183,10 +189,11 @@ def test_response_dataset_gsm8k_with_subset(): # check the first example first_example = dataset.dataset[0] - # only contains messages and task_name - assert len(first_example.keys()) == 2 + # Formatted response rows retain their durable source identity. + assert len(first_example.keys()) == 3 assert "messages" in first_example assert "task_name" in first_example + assert "sample_id" in first_example # check the content assert first_example["messages"][0]["role"] == "user" @@ -243,6 +250,53 @@ def test_gsm8k_subset_defaults_to_main(monkeypatch): assert captured["name"] == "main" +def test_response_dataset_preserves_upstream_uuid(tmp_path: Path): + data_path = tmp_path / "uuid.json" + data_path.write_text( + json.dumps( + [ + { + "input": "question", + "output": "answer", + "uuid": "550e8400-e29b-41d4-a716-446655440000", + } + ] + ), + encoding="utf-8", + ) + + dataset = ResponseDataset(str(data_path)) + + assert dataset.dataset[0]["sample_id"] == ("550e8400-e29b-41d4-a716-446655440000") + + +def test_response_dataset_raw_ordinals_survive_validation_split(tmp_path: Path): + data_path = tmp_path / "split.json" + rows = [ + {"input": f"question-{ordinal}", "output": f"answer-{ordinal}"} + for ordinal in range(8) + ] + data_path.write_text(json.dumps(rows), encoding="utf-8") + + dataset = ResponseDataset(str(data_path), split_validation_size=2, seed=17) + all_rows = [*dataset.dataset, *dataset.val_dataset] + actual_by_question = { + row["messages"][0]["content"]: row["sample_id"] for row in all_rows + } + expected_by_question = { + f"question-{ordinal}": resolve_source_sample_id( + {}, + ordinal, + data_path=str(data_path), + subset=None, + split=None, + ) + for ordinal in range(8) + } + + assert actual_by_question == expected_by_question + + def test_helpsteer3_dataset(): # load the dataset data_config = {"dataset_name": "HelpSteer3"} diff --git a/tests/unit/data/packing/test_lockstep.py b/tests/unit/data/packing/test_lockstep.py new file mode 100644 index 00000000000..a93a99821e2 --- /dev/null +++ b/tests/unit/data/packing/test_lockstep.py @@ -0,0 +1,302 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pickle +from dataclasses import FrozenInstanceError + +import pytest + +from nemo_rl.data.packing import ( + LockstepPackingItem, + SidePackingSpec, + build_lockstep_packing_plan, +) + + +def _items(count: int) -> tuple[LockstepPackingItem, ...]: + return tuple( + LockstepPackingItem(sample_id=f"source-{i}", batch_item_id=i) + for i in range(count) + ) + + +def _round_total_to(multiple: int): + def physical_size(effective_lengths: tuple[int, ...]) -> int: + total = sum(effective_lengths) + return (total + multiple - 1) // multiple * multiple + + return physical_size + + +def test_builds_shared_bins_with_side_local_geometry() -> None: + items = _items(4) + plan = build_lockstep_packing_plan( + batch_uid=9, + items=items, + data_parallel_size=1, + sides=( + SidePackingSpec( + side_id="student", + capacity=16, + raw_lengths=(6, 6, 4, 4), + effective_lengths=(8, 8, 4, 4), + ), + SidePackingSpec( + side_id="teacher", + capacity=16, + raw_lengths=(5, 7, 3, 5), + effective_lengths=(5, 7, 3, 5), + physical_size_fn=_round_total_to(8), + ), + ), + ) + + assert plan.canonical_batch_item_ids == (0, 1, 2, 3) + assert plan.bins == ((0, 1), (3, 2)) + assert plan.sides["student"].rank_bin_indices == ((0, 1),) + assert plan.sides["teacher"].rank_bin_indices == ((0, 1),) + + student = plan.sides["student"] + assert student.raw_cu_seqlens_by_bin == ((0, 6, 12), (0, 4, 8)) + assert student.padded_cu_seqlens_by_bin == ((0, 8, 16), (0, 4, 8)) + assert student.physical_tokens_by_bin == (16, 8) + + teacher = plan.sides["teacher"] + assert teacher.raw_cu_seqlens_by_bin == ((0, 5, 12), (0, 5, 8)) + assert teacher.padded_cu_seqlens_by_bin == ((0, 5, 16), (0, 5, 8)) + assert teacher.physical_tokens_by_bin == (16, 8) + + +def test_item_that_fits_no_existing_bin_starts_the_next_bin() -> None: + plan = build_lockstep_packing_plan( + batch_uid=10, + items=_items(3), + data_parallel_size=1, + sides=( + SidePackingSpec( + side_id="student", + capacity=10, + raw_lengths=(8, 4, 3), + effective_lengths=(8, 4, 3), + ), + SidePackingSpec( + side_id="teacher", + capacity=10, + raw_lengths=(3, 7, 3), + effective_lengths=(3, 7, 3), + ), + ), + ) + + # Item 1 cannot join item 0 on the student side, so it starts bin 2. + # Item 2 then fails bin 1 but fits bin 2, proving that all bins are tried. + assert plan.bins == ((0,), (1, 2)) + + +def test_duplicate_sample_ids_remain_distinct_batch_occurrences() -> None: + items = ( + LockstepPackingItem(sample_id="same-source-row", batch_item_id=100), + LockstepPackingItem(sample_id="same-source-row", batch_item_id=101), + ) + plan = build_lockstep_packing_plan( + batch_uid=3, + items=items, + data_parallel_size=1, + sides=( + SidePackingSpec( + side_id="student", + capacity=16, + raw_lengths=(5, 5), + effective_lengths=(5, 5), + ), + ), + ) + + assert plan.canonical_batch_item_ids == (100, 101) + assert plan.bins == ((100, 101),) + + +def test_plan_is_deterministic_and_immutable() -> None: + items = _items(6) + sides = ( + SidePackingSpec( + side_id="student", + capacity=12, + raw_lengths=(6, 4, 2, 5, 3, 2), + effective_lengths=(6, 4, 2, 5, 3, 2), + ), + SidePackingSpec( + side_id="teacher", + capacity=12, + raw_lengths=(3, 5, 4, 2, 6, 3), + effective_lengths=(3, 5, 4, 2, 6, 3), + ), + ) + + first = build_lockstep_packing_plan( + batch_uid=4, items=items, data_parallel_size=2, sides=sides + ) + second = build_lockstep_packing_plan( + batch_uid=4, items=items, data_parallel_size=2, sides=sides + ) + reversed_sides = build_lockstep_packing_plan( + batch_uid=4, + items=items, + data_parallel_size=2, + sides=tuple(reversed(sides)), + ) + assert first == second + assert first == reversed_sides + assert hash(first) == hash(second) + assert hash(first) == hash(reversed_sides) + assert pickle.loads(pickle.dumps(first)) == first + + with pytest.raises(FrozenInstanceError): + first.batch_uid = 5 # type: ignore[misc] + with pytest.raises(TypeError): + first.sides["student"] = first.sides["student"] # type: ignore[index] + with pytest.raises(FrozenInstanceError): + first.sides["student"].capacity = 20 # type: ignore[misc] + + +def test_rejects_duplicate_batch_item_id() -> None: + with pytest.raises(ValueError, match="batch_item_id must be unique"): + build_lockstep_packing_plan( + batch_uid=0, + items=( + LockstepPackingItem(sample_id="a", batch_item_id=7), + LockstepPackingItem(sample_id="b", batch_item_id=7), + ), + data_parallel_size=1, + sides=( + SidePackingSpec( + side_id="student", + capacity=10, + raw_lengths=(4, 4), + effective_lengths=(4, 4), + ), + ), + ) + + +def test_rejects_effective_length_smaller_than_raw_length() -> None: + with pytest.raises( + ValueError, match="effective length 7 smaller than raw length 8" + ): + SidePackingSpec( + side_id="student", + capacity=10, + raw_lengths=(8,), + effective_lengths=(7,), + ) + + +def test_oversize_error_names_source_sample_and_side() -> None: + with pytest.raises( + ValueError, + match=r"sample_id='raw-row-42'.*batch_item_id=3.*side 'teacher'.*capacity=100", + ): + build_lockstep_packing_plan( + batch_uid=1, + items=(LockstepPackingItem(sample_id="raw-row-42", batch_item_id=3),), + data_parallel_size=1, + sides=( + SidePackingSpec( + side_id="student", + capacity=128, + raw_lengths=(90,), + effective_lengths=(96,), + ), + SidePackingSpec( + side_id="teacher", + capacity=100, + raw_lengths=(97,), + effective_lengths=(104,), + ), + ), + ) + + +def test_gbs48_dp4_equalizes_seven_bins_to_eight_by_shared_split() -> None: + items = _items(48) + student_lengths = (10,) * 12 + (6,) * 12 + (10,) * 12 + (6,) * 12 + teacher_lengths = (6,) * 12 + (10,) * 12 + (10,) * 12 + (7,) * 12 + + initial_counts = [] + for rank in range(4): + start = rank * 12 + stop = start + 12 + rank_plan = build_lockstep_packing_plan( + batch_uid=100 + rank, + items=items[start:stop], + data_parallel_size=1, + sides=( + SidePackingSpec( + side_id="student", + capacity=100, + raw_lengths=student_lengths[start:stop], + effective_lengths=student_lengths[start:stop], + ), + SidePackingSpec( + side_id="teacher", + capacity=100, + raw_lengths=teacher_lengths[start:stop], + effective_lengths=teacher_lengths[start:stop], + ), + ), + ) + initial_counts.append(len(rank_plan.bins)) + assert initial_counts == [2, 2, 2, 1] + assert sum(initial_counts) == 7 + + plan = build_lockstep_packing_plan( + batch_uid=200, + items=items, + data_parallel_size=4, + sides=( + SidePackingSpec( + side_id="student", + capacity=100, + raw_lengths=student_lengths, + effective_lengths=student_lengths, + ), + SidePackingSpec( + side_id="teacher", + capacity=100, + raw_lengths=teacher_lengths, + effective_lengths=teacher_lengths, + ), + ), + ) + + assert plan.bins == ( + tuple(range(0, 10)), + (10, 11), + tuple(range(12, 22)), + (22, 23), + tuple(range(24, 34)), + (34, 35), + tuple(range(36, 47)), + (47,), + ) + expected_routing = ((0, 1), (2, 3), (4, 5), (6, 7)) + assert plan.sides["student"].rank_bin_indices == expected_routing + assert plan.sides["teacher"].rank_bin_indices == expected_routing + assert len(plan.bins) == 8 + assert plan.equalization_splits == 1 + assert sorted(item_id for bin_items in plan.bins for item_id in bin_items) == list( + range(48) + ) + for side in plan.sides.values(): + assert all(tokens <= side.capacity for tokens in side.physical_tokens_by_bin) diff --git a/tests/unit/data/test_cross_tokenizer_collate.py b/tests/unit/data/test_cross_tokenizer_collate.py index 77d38253a25..ba939dce00b 100644 --- a/tests/unit/data/test_cross_tokenizer_collate.py +++ b/tests/unit/data/test_cross_tokenizer_collate.py @@ -404,7 +404,8 @@ def test_long_text_truncates_trailing_tokens_not_drops_sample(self): # Captured from the collator before the optional overlength and CE-turn # behaviors existed; the default batch must stay bitwise identical. Its first # row fills each side's context, so padding to the longest row keeps the -# full-context width. +# full-context width. ``alignment_0_pair_valid`` was recaptured when pairs +# over padding became invalid. _TEXT_SNAPSHOT_CHECKSUMS = { "input_ids": 3661199532, "input_lengths": 868467641, @@ -414,7 +415,7 @@ def test_long_text_truncates_trailing_tokens_not_drops_sample(self): "teacher_0_input_ids": 548984487, "teacher_0_input_lengths": 137384594, "teacher_0_token_mask": 3787835366, - "alignment_0_pair_valid": 64817739, + "alignment_0_pair_valid": 2277715625, "alignment_0_pair_is_correct": 13516310, "alignment_0_student_chunk_id": 1888554585, "alignment_0_teacher_chunk_id": 1481747482, @@ -883,6 +884,29 @@ def test_all_empty_batch_pads_to_the_divisor(self): assert out["alignment_0_student_chunk_id"].eq(-1).all() assert out["alignment_0_teacher_chunk_id"].eq(-1).all() + def test_pair_valid_counts_real_chunks(self): + """``pair_valid`` counts the pairs holding real tokens, including the + student-only start-of-text pair, and none of the pairs over padding. + """ + batch = [_datum(text, idx=i) for i, text in enumerate(self._TEXTS)] + out = _text_collator( + ctx_length_student=16, + ctx_length_teachers=[64, 16], + student=_word_tokenizer(bos=True), + add_special_tokens_in_text=True, + )(batch) + real = 0 + for row in range(len(batch)): + student = out["alignment_0_student_chunk_id"][row] + teacher = out["alignment_0_teacher_chunk_id"][row] + real += len( + set(student[student >= 0].tolist()) + | set(teacher[teacher >= 0].tolist()) + ) + # One pair per word plus each row's start-of-text pair. + assert real == 8 + 3 + assert out["alignment_0_pair_valid"].sum().item() == real + class TestCollatorPadTokenFallback: def test_missing_pad_token_set_from_eos(self): diff --git a/tests/unit/data/test_csd_collate.py b/tests/unit/data/test_csd_collate.py new file mode 100644 index 00000000000..2a1074399d8 --- /dev/null +++ b/tests/unit/data/test_csd_collate.py @@ -0,0 +1,688 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Hermetic CPU tests for ``CSDCrossTokenizerCollator``. + +CSD batches are built by the shared ``CrossTokenizerCollator`` (tested in +``test_cross_tokenizer_collate.py``, ``test_cross_tokenizer_chat_eot.py`` and +``test_cross_tokenizer_native_chat.py``). These tests cover what CSD adds, +row ``sample_id`` values and MoE router-replay routes, and the batches the +shared collator builds with the options CSD's trainer turns on: special tokens +and overlength clipping in text mode, and whole-final-turn CE in chat mode. +Tokenizers are small local fast tokenizers; nothing is downloaded. +""" + +from __future__ import annotations + +from dataclasses import fields as dataclass_fields +from typing import Any, Optional +from unittest.mock import patch + +import pytest +import torch +from tokenizers import Tokenizer +from tokenizers.models import BPE, WordLevel +from tokenizers.pre_tokenizers import WhitespaceSplit +from tokenizers.processors import TemplateProcessing +from transformers import PreTrainedTokenizerFast + +from nemo_rl.algorithms.x_token.token_aligner import AlignmentBatch, TokenAligner +from nemo_rl.data.cross_tokenizer.cross_tokenizer_collate import ( + CrossTokenizerCollator, + CrossTokenizerCollatorConfig, +) +from nemo_rl.data.cross_tokenizer.csd_collate import CSDCrossTokenizerCollator +from nemo_rl.data.interfaces import TaskDataSpec +from nemo_rl.data.processors import kd_data_processor +from tests.unit.data.test_cross_tokenizer_collate import _assert_batches_equal +from tests.unit.data.test_cross_tokenizer_native_chat import _messages, _native_pair + +_TEXT = CrossTokenizerCollatorConfig(mode="text") +_CHAT = CrossTokenizerCollatorConfig(mode="chat") +_NATIVE_CHAT = CrossTokenizerCollatorConfig( + mode="chat", include_thinking_in_loss=True, native_thinking_alignment=True +) + +# --------------------------------------------------------------------------- +# Collator construction, mirroring the trainer's CSD options. +# --------------------------------------------------------------------------- + + +def _shared_options(config: CrossTokenizerCollatorConfig) -> dict[str, bool]: + """Shared-collator options CSD's trainer turns on for ``config``.""" + options = {"reject_overlength_text": True} + if config.mode == "text": + options["add_special_tokens_in_text"] = True + return options + + +def _collator_args( + config: CrossTokenizerCollatorConfig, + student: Any, + teachers: list[Optional[Any]], + *, + ctx_length_student: int = 64, + ctx_length_teachers: Optional[list[int]] = None, + make_seq_div_by_student: int = 1, + make_seq_div_by_teachers: Optional[list[int]] = None, +) -> dict[str, Any]: + """Constructor arguments common to the CSD and the shared collator. + + A ``None`` entry in ``teachers`` is a same-tokenizer teacher; every other + entry is a cross-tokenizer teacher's tokenizer. + """ + n = len(teachers) + return { + "config": config, + "student_tokenizer": student, + "teacher_tokenizers": [student if t is None else t for t in teachers], + "aligners": [ + None + if t is None + else TokenAligner(student, t, projection_matrix_path="unused.pt") + for t in teachers + ], + "ctx_length_student": ctx_length_student, + "ctx_length_teachers": ctx_length_teachers or [64] * n, + "drop_first_assistant_chunk_kl_by_teacher": [False] * n, + "make_seq_div_by_student": make_seq_div_by_student, + "make_seq_div_by_teachers": make_seq_div_by_teachers or [1] * n, + } + + +def _csd_collator( + config: CrossTokenizerCollatorConfig, + student: Any, + teachers: list[Optional[Any]], + *, + require_routed_experts: bool = False, + **sizes: Any, +) -> CSDCrossTokenizerCollator: + """CSD collator built with the options the trainer passes for ``config``.""" + return CSDCrossTokenizerCollator( + **_collator_args(config, student, teachers, **sizes), + require_routed_experts=require_routed_experts, + clip_overlength_text=config.mode == "text", + **_shared_options(config), + ) + + +def _shared_collator( + config: CrossTokenizerCollatorConfig, + student: Any, + teachers: list[Optional[Any]], + **sizes: Any, +) -> CrossTokenizerCollator: + """Shared collator with the same arguments and options as ``_csd_collator``.""" + return CrossTokenizerCollator( + **_collator_args(config, student, teachers, **sizes), + **_shared_options(config), + ) + + +# --------------------------------------------------------------------------- +# Local fast tokenizers +# --------------------------------------------------------------------------- + +_WORDS = ("alpha", "beta", "gamma", "delta", "epsilon") + + +def _text_tokenizer(*, bos: str, char_level: bool) -> PreTrainedTokenizerFast: + """Fast tokenizer whose special tokens put ``bos`` at the start of each row. + + Word-level tokenizers emit one token per whitespace-separated word; + character-level ones one token per non-space character. + """ + pieces = sorted(set("".join(_WORDS))) if char_level else list(_WORDS) + vocab = { + token: i for i, token in enumerate(("[UNK]", "[PAD]", "", bos, *pieces)) + } + backend = Tokenizer( + BPE(vocab=vocab, merges=[], unk_token="[UNK]") + if char_level + else WordLevel(vocab, unk_token="[UNK]") + ) + backend.pre_tokenizer = WhitespaceSplit() + backend.post_processor = TemplateProcessing( + single=f"{bos} $A", special_tokens=[(bos, vocab[bos])] + ) + return PreTrainedTokenizerFast( + tokenizer_object=backend, + unk_token="[UNK]", + pad_token="[PAD]", + eos_token="", + bos_token=bos, + ) + + +def _text_pair() -> tuple[PreTrainedTokenizerFast, PreTrainedTokenizerFast]: + """Word-level student marked ```` and character-level teacher marked ````.""" + return ( + _text_tokenizer(bos="", char_level=False), + _text_tokenizer(bos="", char_level=True), + ) + + +_CHAT_SPECIALS = ("", "", "", "", "") + + +def _chat_tokenizer( + *, teacher: bool = False, suffix: str = "" +) -> PreTrainedTokenizerFast: + """Word-level fast tokenizer with a chat template; role markers and EOT are special. + + Each message renders as its role marker, its content, then ``suffix``. + The student marks roles ````/````, the teacher ````/````. + """ + tokens = ( + "[UNK]", + "[PAD]", + "", + *_CHAT_SPECIALS, + "Question", + "Answer", + "Again", + "Done", + ) + backend = Tokenizer( + WordLevel({token: i for i, token in enumerate(tokens)}, unk_token="[UNK]") + ) + backend.pre_tokenizer = WhitespaceSplit() + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=backend, + unk_token="[UNK]", + pad_token="[PAD]", + eos_token="", + additional_special_tokens=list(_CHAT_SPECIALS), + ) + user, assistant = ("", "") if teacher else ("", "") + tokenizer.chat_template = ( + "{% for message in messages %}" + f"{{{{ ({user!r} if message['role'] == 'user' else {assistant!r})" + f" + message['content'] + {suffix!r} }}}}" + "{% endfor %}" + ) + return tokenizer + + +# --------------------------------------------------------------------------- +# Rows, as kd_data_processor emits them +# --------------------------------------------------------------------------- + +_SINGLE_TURN = [ + {"role": "user", "content": "Question"}, + {"role": "assistant", "content": "Answer"}, +] +_MULTI_TURN = [ + *_SINGLE_TURN, + {"role": "user", "content": "Again"}, + {"role": "assistant", "content": "Done"}, +] + + +def _row( + messages: list[dict[str, Any]], + idx: int = 0, + *, + sample_id: Any = None, + message_loss_mask: Optional[list[int]] = None, +) -> dict[str, Any]: + """``kd_data_processor`` output; ``sample_id`` defaults to ``doc-{idx}``.""" + return kd_data_processor( + { + "messages": messages, + "message_loss_mask": message_loss_mask, + "sample_id": f"doc-{idx}" if sample_id is None else sample_id, + }, + TaskDataSpec(), + None, + None, + idx, + ) + + +def _text_row(text: str, idx: int = 0, **kwargs: Any) -> dict[str, Any]: + return _row([{"role": "assistant", "content": text}], idx, **kwargs) + + +def _supervised( + tokenizer: PreTrainedTokenizerFast, + batch: Any, + mask_key: str, + *, + ids_key: str = "input_ids", + row: int = 0, +) -> list[str]: + """Tokens of ``batch[ids_key][row]`` selected by ``batch[mask_key][row]``.""" + ids = batch[ids_key][row][batch[mask_key][row].bool()] + return tokenizer.convert_ids_to_tokens(ids.tolist()) + + +_ALIGNMENT_KEYS = {f"alignment_0_{f.name}" for f in dataclass_fields(AlignmentBatch)} +_STUDENT_KEYS = { + "input_ids", + "input_lengths", + "token_mask", + "sample_mask", + "idx", + "sample_id", +} +_TEACHER_0_KEYS = { + "teacher_0_input_ids", + "teacher_0_input_lengths", + "teacher_0_token_mask", + *_ALIGNMENT_KEYS, +} + +# --------------------------------------------------------------------------- +# Output contract +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("mode", ["text", "chat"]) +def test_batch_is_the_shared_batch_plus_sample_ids(mode: str): + # Teacher 0 is cross-tokenizer; teacher 1 shares the student tokenizer + # and emits nothing extra. + if mode == "text": + config = _TEXT + student, teacher = _text_pair() + batch = [ + _text_row("alpha beta gamma", 0, sample_id="doc-b"), + _text_row("delta", 1, sample_id=0), + ] + else: + config = _CHAT + student = _chat_tokenizer() + teacher = _chat_tokenizer(teacher=True, suffix="\n") + batch = [ + _row(_MULTI_TURN, 0, sample_id="doc-b"), + _row(_SINGLE_TURN, 1, sample_id=0), + ] + + out = _csd_collator(config, student, [teacher, None], make_seq_div_by_student=8)( + batch + ) + + assert set(out.keys()) == _STUDENT_KEYS | _TEACHER_0_KEYS + # Row identities keep their values (an integer 0 is a valid id) and order. + assert out["sample_id"] == ["doc-b", 0] + shared = _shared_collator( + config, student, [teacher, None], make_seq_div_by_student=8 + )(batch) + _assert_batches_equal({k: v for k, v in out.items() if k != "sample_id"}, shared) + + +def test_text_token_mask_is_the_attention_mask(): + student, teacher = _text_pair() + out = _csd_collator(_TEXT, student, [teacher], make_seq_div_by_student=8)( + [_text_row("alpha beta gamma", 0), _text_row("delta", 1)] + ) + for prefix in ("", "teacher_0_"): + width = out[f"{prefix}input_ids"].shape[1] + attention = ( + torch.arange(width)[None, :] < out[f"{prefix}input_lengths"][:, None] + ) + assert torch.equal(out[f"{prefix}token_mask"], attention.long()), prefix + assert "kd_token_mask" not in out + + +def test_drop_first_flags_must_match_teacher_count(): + student, teacher = _text_pair() + args = _collator_args(_TEXT, student, [teacher]) + args["drop_first_assistant_chunk_kl_by_teacher"] = [] + with pytest.raises(ValueError, match="drop_first_assistant_chunk_kl_by_teacher"): + CSDCrossTokenizerCollator( + **args, require_routed_experts=False, **_shared_options(_TEXT) + ) + + +# --------------------------------------------------------------------------- +# sample_id +# --------------------------------------------------------------------------- + +_MISSING = object() + + +@pytest.mark.parametrize( + "sample_id", + [_MISSING, None, "", " \t\n"], + ids=["missing", "none", "empty", "whitespace"], +) +def test_rows_without_a_sample_id_are_rejected_before_tokenization(sample_id): + student, teacher = _text_pair() + collator = _csd_collator(_TEXT, student, [teacher]) + batch = [_text_row("beta", 0), _text_row("delta", 1)] + if sample_id is _MISSING: + del batch[1]["sample_id"] + else: + batch[1]["sample_id"] = sample_id + with ( + patch.object(CrossTokenizerCollator, "__call__") as shared_call, + pytest.raises( + ValueError, match=r"non-empty sample_id on every row.*\(batch row 1\)" + ), + ): + collator(batch) + shared_call.assert_not_called() + + +# --------------------------------------------------------------------------- +# Router replay +# --------------------------------------------------------------------------- + + +def _router_replay_collator( + config: CrossTokenizerCollatorConfig, student: PreTrainedTokenizerFast +) -> CSDCrossTokenizerCollator: + return _csd_collator( + config, + student, + [None], + make_seq_div_by_student=8, + require_routed_experts=True, + ) + + +def test_router_replay_pads_rollout_routes_to_the_student_width(): + student, _ = _text_pair() + text = "alpha beta gamma" + # The rollout recorded the student tokenization, including its . + token_ids = student(text)["input_ids"] + routes = torch.arange(len(token_ids) * 2 * 2).reshape(len(token_ids), 2, 2) + row = _text_row(text) + # Lists exercise the JSON-serialized rollout representation. + row["message_log"][0]["token_ids"] = token_ids + row["message_log"][0]["routed_experts"] = routes.tolist() + + out = _router_replay_collator(_TEXT, student)([row]) + + assert out["input_ids"][0, : len(token_ids)].tolist() == token_ids + assert out["routed_experts"].shape == (1, 8, 2, 2) + assert torch.equal(out["routed_experts"][0, : len(token_ids)], routes) + assert torch.count_nonzero(out["routed_experts"][0, len(token_ids) :]) == 0 + + +def test_router_replay_fails_clearly_when_routes_are_unavailable(): + student, _ = _text_pair() + row = _text_row("alpha beta") + row["message_log"][0]["token_ids"] = student("alpha beta")["input_ids"] + with pytest.raises( + RuntimeError, + match=r"policy\.router_replay\.enabled=true.*missing routed_experts", + ): + _router_replay_collator(_TEXT, student)([row]) + + +def test_router_replay_rejects_routes_for_a_different_student_tokenization(): + student, _ = _text_pair() + text = "alpha beta gamma" + row = _text_row(text) + # Recorded without the student's , so positions would be off by one. + token_ids = student(text, add_special_tokens=False)["input_ids"] + row["message_log"][0]["token_ids"] = token_ids + row["message_log"][0]["routed_experts"] = torch.zeros(len(token_ids), 2, 2) + with pytest.raises(RuntimeError, match="do not match.*student tokenization"): + _router_replay_collator(_TEXT, student)([row]) + + +def test_router_replay_chat_concatenates_per_message_routes(): + student = _chat_tokenizer() + messages: list[dict[str, Any]] = [dict(message) for message in _SINGLE_TURN] + rendered = student.apply_chat_template(messages, tokenize=False) + token_ids = student(rendered, add_special_tokens=False)["input_ids"] + # Question | Answer + split = 3 + routes = torch.arange(len(token_ids) * 2 * 2).reshape(len(token_ids), 2, 2) + for message, part in zip(messages, (slice(None, split), slice(split, None))): + message["token_ids"] = token_ids[part] + message["routed_experts"] = routes[part] + + out = _router_replay_collator(_CHAT, student)([_row(messages)]) + + assert out["routed_experts"].shape == (1, 8, 2, 2) + assert torch.equal(out["routed_experts"][0, : len(token_ids)], routes) + assert torch.count_nonzero(out["routed_experts"][0, len(token_ids) :]) == 0 + # The supervised content and EOT replay the final message's routes after + # its header. + supervised = out["token_mask"][0].bool() + assert torch.equal(out["routed_experts"][0, supervised], routes[split + 1 :]) + + +# --------------------------------------------------------------------------- +# Text mode: special tokens, pad-to-longest, overlength clipping +# --------------------------------------------------------------------------- + + +def test_text_rows_carry_special_tokens_and_pad_to_the_longest_row(): + student, teacher = _text_pair() + texts = ("alpha beta gamma", "delta") + out = _csd_collator( + _TEXT, + student, + [teacher], + make_seq_div_by_student=8, + make_seq_div_by_teachers=[4], + )([_text_row(text, i) for i, text in enumerate(texts)]) + + # Lengths count each tokenizer's marker: 1 + 3 words and 1 + 14 characters + # for the longest row. Widths round the longest row up to the divisor, + # not to the 64-token contexts. + for prefix, tokenizer, lengths, width in ( + ("", student, [4, 2], 8), + ("teacher_0_", teacher, [15, 6], 16), + ): + ids = out[f"{prefix}input_ids"] + assert out[f"{prefix}input_lengths"].tolist() == lengths, prefix + assert ids.shape == (2, width), prefix + assert ids[:, 0].eq(tokenizer.bos_token_id).all(), prefix + for row, (text, n) in enumerate(zip(texts, lengths)): + assert ids[row, :n].tolist() == tokenizer(text)["input_ids"], prefix + assert ids[row, n:].eq(tokenizer.pad_token_id).all(), prefix + + +@pytest.mark.parametrize( + "sizes,text,prefix", + [ + pytest.param( + {"ctx_length_student": 4}, + "alpha beta gamma delta", + "alpha beta gamma ", + id="student", + ), + pytest.param( + {"ctx_length_student": 6, "make_seq_div_by_student": 4}, + "alpha beta gamma delta", + "alpha beta gamma ", + id="student-after-rounding", + ), + pytest.param( + {"ctx_length_teachers": [10, 64]}, + "alpha beta gamma", + "alpha beta ", + id="cross-tokenizer-teacher", + ), + pytest.param( + {"ctx_length_teachers": [64, 4]}, + "alpha beta gamma delta", + "alpha beta gamma ", + id="same-tokenizer-teacher", + ), + pytest.param( + {"ctx_length_teachers": [64, 6], "make_seq_div_by_teachers": [1, 4]}, + "alpha beta gamma delta", + "alpha beta gamma ", + id="same-tokenizer-teacher-after-rounding", + ), + ], +) +def test_text_overlength_rows_are_clipped_to_a_shared_prefix(sizes, text, prefix): + # Token counts include each tokenizer's marker: "alpha beta gamma delta" + # is 5 student tokens, "alpha beta gamma" 15 teacher tokens. Teacher 1 + # shares the student tokenization but has its own context and divisor. + student, teacher = _text_pair() + teachers = [teacher, None] + out = _csd_collator(_TEXT, student, teachers, **sizes)( + [_text_row("beta", 0), _text_row(text, 1)] + ) + + # The clipped batch is the one the shared collator builds from the prefix, + # for the student and every teacher, and one more character does not fit. + shared = _shared_collator(_TEXT, student, teachers, **sizes) + expected = shared([_text_row("beta", 0), _text_row(prefix, 1)]) + _assert_batches_equal({k: v for k, v in out.items() if k != "sample_id"}, expected) + with pytest.raises(ValueError, match="overlength rows are rejected$"): + shared([_text_row(text[: len(prefix) + 1], 1)]) + + +def test_text_row_without_a_fitting_prefix_raises(): + student, teacher = _text_pair() + # Even one character tokenizes to the marker plus one token. + collator = _csd_collator(_TEXT, student, [teacher], ctx_length_student=1) + with pytest.raises( + ValueError, + match=r"^sample idx=0: no non-empty prefix of the text fits every " + r"context length$", + ): + collator([_text_row("beta", 0)]) + + +def test_clip_overlength_text_requires_text_mode(): + student, teacher = _chat_pair() + with pytest.raises(ValueError, match="clip_overlength_text requires mode='text'"): + CSDCrossTokenizerCollator( + **_collator_args(_CHAT, student, [teacher]), + require_routed_experts=False, + clip_overlength_text=True, + ) + + +# --------------------------------------------------------------------------- +# Chat mode: CE and KD on the final assistant turn's content and EOT +# --------------------------------------------------------------------------- + + +def _chat_pair() -> tuple[PreTrainedTokenizerFast, PreTrainedTokenizerFast]: + return _chat_tokenizer(), _chat_tokenizer(teacher=True, suffix="\n") + + +def _assert_kd_alignment_matches_masks(out: Any) -> None: + """Aligned positions are exactly the KD positions on each side.""" + for side, mask_key in ( + ("student", "token_mask"), + ("teacher", "teacher_0_token_mask"), + ): + chunks = out[f"alignment_0_{side}_chunk_id"][0] + mask = out[mask_key][0] + assert chunks[mask == 1].ge(0).all(), side + assert chunks[mask == 0].eq(-1).all(), side + + +def test_chat_single_turn_masks(): + student, teacher = _chat_pair() + out = _csd_collator(_CHAT, student, [teacher], make_seq_div_by_student=8)( + [_row(_SINGLE_TURN)] + ) + assert _supervised(student, out, "token_mask") == ["Answer", ""] + assert "kd_token_mask" not in out + assert _supervised( + teacher, out, "teacher_0_token_mask", ids_key="teacher_0_input_ids" + ) == ["Answer", ""] + _assert_kd_alignment_matches_masks(out) + + +@pytest.mark.parametrize( + "message_loss_mask", + [None, [0, 0, 0, 1]], + ids=["earlier-and-final-selected", "final-selected"], +) +def test_chat_multi_turn_supervises_only_the_final_turn(message_loss_mask): + student, teacher = _chat_pair() + row = _row(_MULTI_TURN, message_loss_mask=message_loss_mask) + out = _csd_collator(_CHAT, student, [teacher])([row]) + + # Earlier turns stay in the input as context. + assert student.convert_ids_to_tokens(out["input_ids"][0].tolist()) == [ + "", + "Question", + "", + "", + "Answer", + "", + "", + "Again", + "", + "", + "Done", + "", + ] + assert _supervised(student, out, "token_mask") == ["Done", ""] + assert _supervised( + teacher, out, "teacher_0_token_mask", ids_key="teacher_0_input_ids" + ) == ["Done", ""] + # The earlier assistant turn gets neither CE nor KD. + assert not out["token_mask"][0, 3:6].any() + _assert_kd_alignment_matches_masks(out) + + +@pytest.mark.parametrize( + "message_loss_mask", + [[0, 0, 0, 0], [0, 1, 0, 0]], + ids=["nothing-selected", "earlier-turn-selected"], +) +def test_chat_selection_without_the_final_turn_is_rejected(message_loss_mask): + student, teacher = _chat_pair() + collator = _csd_collator(_CHAT, student, [teacher]) + with pytest.raises( + ValueError, + match=r"^student, sample idx=0: The final assistant message must be " + r"selected for supervision$", + ): + collator([_row(_MULTI_TURN, message_loss_mask=message_loss_mask)]) + + +def test_chat_same_tokenizer_teacher_bypasses_alignment(): + student = _chat_tokenizer() + out = _csd_collator(_CHAT, student, [None])([_row(_SINGLE_TURN)]) + assert set(out.keys()) == _STUDENT_KEYS + assert _supervised(student, out, "token_mask") == ["Answer", ""] + + +def test_chat_template_with_a_footer_supervises_content_and_eot(): + # A footer after the last message is context, not supervision. + student = _chat_tokenizer() + student.chat_template += "{{ '' }}" + out = _csd_collator(_CHAT, student, [None])([_row(_SINGLE_TURN)]) + assert _supervised(student, out, "token_mask") == ["Answer", ""] + + +# --------------------------------------------------------------------------- +# Native thinking alignment +# --------------------------------------------------------------------------- + + +def test_native_thinking_uses_the_shared_native_path_and_adds_sample_ids(): + student, teacher, _ = _native_pair() + sizes = {"ctx_length_student": 256, "ctx_length_teachers": [256]} + row = { + "idx": 7, + "message_loss_mask": [0, 1], + "message_log": _messages(), + "loss_multiplier": 1.0, + "sample_id": "doc-7", + } + out = _csd_collator(_NATIVE_CHAT, student, [teacher], **sizes)([row]) + shared = _shared_collator(_NATIVE_CHAT, student, [teacher], **sizes)([row]) + + assert out["sample_id"] == ["doc-7"] + assert "kd_token_mask" not in out + _assert_batches_equal({k: v for k, v in out.items() if k != "sample_id"}, shared) diff --git a/tests/unit/data/test_data_config.py b/tests/unit/data/test_data_config.py new file mode 100644 index 00000000000..57ace4ea788 --- /dev/null +++ b/tests/unit/data/test_data_config.py @@ -0,0 +1,45 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pydantic import TypeAdapter + +from nemo_rl.data import DataConfig + + +def test_data_config_preserves_native_openai_chat_adapter_fields() -> None: + train = { + "dataset_name": "openai_format", + "data_path": "/openmath-chat.jsonl", + "data_files": None, + "text_key": None, + "characters_per_sample": None, + "processor": "kd_data_processor", + "split": "train", + "chat_key": "messages", + "use_preserving_dataset": False, + "system_key": None, + "system_prompt": None, + "tool_key": None, + "prompt_file": None, + } + + validated = TypeAdapter(DataConfig).validate_python( + { + "max_input_seq_length": 4096, + "shuffle": False, + "train": train, + } + ) + + assert validated["train"] == train diff --git a/tests/unit/data/test_data_processor.py b/tests/unit/data/test_data_processor.py index 071a2390535..95578698654 100644 --- a/tests/unit/data/test_data_processor.py +++ b/tests/unit/data/test_data_processor.py @@ -591,7 +591,8 @@ def test_registry_resolves_kd_data_processor(self): def test_output_keys_and_values(self): out = kd_data_processor( datum_dict={ - "messages": [{"role": "assistant", "content": "the quick brown fox"}] + "messages": [{"role": "assistant", "content": "the quick brown fox"}], + "sample_id": "corpus#3", }, task_data_spec=self._spec(), tokenizer=DummyTokenizer(), # must not be called @@ -606,12 +607,14 @@ def test_output_keys_and_values(self): assert out["extra_env_info"] is None assert out["loss_multiplier"] == 1.0 assert out["idx"] == 3 + assert out["sample_id"] == "corpus#3" def test_task_name_forwarded_when_present(self): out = kd_data_processor( datum_dict={ "messages": [{"role": "assistant", "content": "hello"}], "task_name": "code", + "sample_id": "corpus#0", }, task_data_spec=self._spec(), tokenizer=DummyTokenizer(), @@ -622,7 +625,10 @@ def test_task_name_forwarded_when_present(self): def test_task_name_absent_when_not_in_datum(self): out = kd_data_processor( - datum_dict={"messages": [{"role": "assistant", "content": "hello"}]}, + datum_dict={ + "messages": [{"role": "assistant", "content": "hello"}], + "sample_id": "corpus#0", + }, task_data_spec=self._spec(), tokenizer=DummyTokenizer(), max_seq_length=128, @@ -635,7 +641,10 @@ def test_does_not_emit_collator_keys(self): # processor's. If a future change emits any of these keys, the # CrossTokenizerCollator's contract is broken. out = kd_data_processor( - datum_dict={"messages": [{"role": "assistant", "content": "hello"}]}, + datum_dict={ + "messages": [{"role": "assistant", "content": "hello"}], + "sample_id": "corpus#0", + }, task_data_spec=self._spec(), tokenizer=DummyTokenizer(), max_seq_length=128, @@ -662,7 +671,10 @@ def encode(self, *a, **kw): # noqa: ARG002 # Should not raise — the processor must not touch the tokenizer. _ = kd_data_processor( - datum_dict={"messages": [{"role": "assistant", "content": "hello"}]}, + datum_dict={ + "messages": [{"role": "assistant", "content": "hello"}], + "sample_id": "corpus#0", + }, task_data_spec=self._spec(), tokenizer=StrictTokenizer(), max_seq_length=128, @@ -676,7 +688,10 @@ def test_long_text_not_truncated_in_processor(self): # max_seq_length. long_text = "a" * 10_000 out = kd_data_processor( - datum_dict={"messages": [{"role": "assistant", "content": long_text}]}, + datum_dict={ + "messages": [{"role": "assistant", "content": long_text}], + "sample_id": "corpus#0", + }, task_data_spec=self._spec(), tokenizer=DummyTokenizer(), max_seq_length=128, @@ -685,3 +700,34 @@ def test_long_text_not_truncated_in_processor(self): assert out["message_log"][0]["content"] == long_text # length is a fake placeholder for the kd pipeline. assert out["length"] == 0 + + def test_preserves_identity_tools_and_deep_copies_nested_messages(self): + source = { + "messages": [ + { + "role": "assistant", + "content": "calling", + "tool_calls": [ + {"function": {"name": "weather", "arguments": {"city": "SF"}}} + ], + } + ], + "tools": [{"type": "function", "function": {"name": "weather"}}], + "sample_id": "corpus#9", + } + out = kd_data_processor( + datum_dict=source, + task_data_spec=self._spec(), + tokenizer=DummyTokenizer(), + max_seq_length=128, + idx=4, + ) + assert out["sample_id"] == "corpus#9" + assert out["tools"] == source["tools"] + out["message_log"][0]["tool_calls"][0]["function"]["arguments"]["city"] = "NYC" + out["tools"][0]["function"]["name"] = "changed" + assert ( + source["messages"][0]["tool_calls"][0]["function"]["arguments"]["city"] + == "SF" + ) + assert source["tools"][0]["function"]["name"] == "weather" diff --git a/tests/unit/distributed/test_batched_data_dict.py b/tests/unit/distributed/test_batched_data_dict.py index 60358a6cc28..8bee4131a0c 100644 --- a/tests/unit/distributed/test_batched_data_dict.py +++ b/tests/unit/distributed/test_batched_data_dict.py @@ -16,6 +16,11 @@ import torch from nemo_rl.data.multimodal_utils import PackedTensor +from nemo_rl.data.packing import ( + LockstepPackingItem, + SidePackingSpec, + build_lockstep_packing_plan, +) from nemo_rl.data_plane.codec import pad_field from nemo_rl.data_plane.schema import ( OPD_FULL_HIDDEN_STATES_FIELD, @@ -29,6 +34,88 @@ ) +def test_shard_by_packing_plan_follows_shared_ids_and_side_geometry(): + plan = build_lockstep_packing_plan( + batch_uid=12, + items=tuple( + LockstepPackingItem(sample_id=f"row-{i}", batch_item_id=100 + i) + for i in range(8) + ), + data_parallel_size=2, + sides=( + SidePackingSpec( + side_id="student", + capacity=10, + raw_lengths=(4, 4, 2, 2, 3, 3, 2, 2), + effective_lengths=(4, 4, 2, 2, 3, 3, 2, 2), + ), + SidePackingSpec( + side_id="teacher_0", + capacity=12, + raw_lengths=(3, 3, 3, 3, 4, 4, 2, 2), + effective_lengths=(3, 3, 3, 3, 4, 4, 2, 2), + ), + ), + ) + # Present the physical batch in a deliberately non-canonical order. The + # explicit IDs, rather than row positions, drive routing. + order = [3, 0, 2, 1, 7, 5, 4, 6] + batch = BatchedDataDict( + { + "batch_item_id": [100 + i for i in order], + "value": torch.tensor(order), + "label": [f"v{i}" for i in order], + } + ) + shards = batch.shard_by_packing_plan(plan, side_id="teacher_0") + assert len(shards) == 2 + for rank, shard in enumerate(shards): + bin_indices = plan.sides["teacher_0"].rank_bin_indices[rank] + expected_ids = [ + item_id for bin_index in bin_indices for item_id in plan.bins[bin_index] + ] + assert shard["batch_item_id"] == expected_ids + assert shard["value"].tolist() == [item_id - 100 for item_id in expected_ids] + assert shard.elem_counts_per_gb == [4] + assert shard.lockstep_packing_plan is plan + assert shard.lockstep_side_id == "teacher_0" + assert shard.lockstep_bin_indices == bin_indices + assert shard.micro_batch_lengths == [ + [ + plan.sides["teacher_0"].physical_tokens_by_bin[bin_index] + for bin_index in bin_indices + ] + ] + global_batch = shard.get_batch(0) + assert global_batch.lockstep_packing_plan is plan + assert global_batch.lockstep_side_id == "teacher_0" + for ( + microbatch + ) in global_batch.make_microbatch_iterator_for_packable_sequences(): + assert microbatch.lockstep_packing_plan is plan + assert microbatch.lockstep_side_id == "teacher_0" + + +def test_shard_by_packing_plan_rejects_missing_occurrence_id(): + plan = build_lockstep_packing_plan( + batch_uid=13, + items=(LockstepPackingItem(sample_id="row", batch_item_id=7),), + data_parallel_size=1, + sides=( + SidePackingSpec( + side_id="student", + capacity=8, + raw_lengths=(4,), + effective_lengths=(4,), + ), + ), + ) + with pytest.raises(ValueError, match="requires a batch_item_id"): + BatchedDataDict({"value": torch.tensor([1])}).shard_by_packing_plan( + plan, side_id="student" + ) + + def test_shard_by_batch_size_basic(): """Test basic functionality of shard_by_batch_size with tensor data.""" # Create a sample batch with tensor data diff --git a/tests/unit/distributed/test_model_utils.py b/tests/unit/distributed/test_model_utils.py index c9a728417ac..ada90a37d00 100644 --- a/tests/unit/distributed/test_model_utils.py +++ b/tests/unit/distributed/test_model_utils.py @@ -17,6 +17,9 @@ import pytest import ray import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F from nemo_rl.algorithms.logits_sampling_utils import apply_top_k_top_p from nemo_rl.distributed.model_utils import ( @@ -32,7 +35,10 @@ _compute_distributed_selected_logprobs, _get_tokens_on_this_cp_rank, allgather_cp_sharded_tensor, + cp_load_balanced_segments, + distributed_selected_logprobs, distributed_vocab_topk, + distributed_vocab_topk_logz_force, from_parallel_logits_to_logprobs, from_parallel_logits_to_logprobs_packed_sequences, gather_logits_at_global_indices, @@ -1997,3 +2003,428 @@ def test_student_teacher_kernels_normalize_bf16_logits_in_fp32( expected_kl, _, _ = _student_teacher_reference(student, teacher) torch.testing.assert_close(reverse_kl, expected_kl, rtol=1e-5, atol=1e-5) + + +# ── Vocab-parallel selected log-probs and sparse top-k primitives ─────────── + + +@pytest.mark.parametrize("cp_size", [1, 2, 4]) +def test_cp_load_balanced_segments_match_cp_rank_tokens(cp_size): + """The segments describe exactly the rows ``_get_tokens_on_this_cp_rank`` keeps.""" + full_seq_len = 5 * 2 * cp_size + tokens = torch.arange(full_seq_len).unsqueeze(0) + for cp_rank in range(cp_size): + reconstructed = [] + expected_local_start = 0 + for local_start, global_start, length in cp_load_balanced_segments( + full_seq_len, cp_rank, cp_size + ): + assert local_start == expected_local_start + reconstructed.extend(range(global_start, global_start + length)) + expected_local_start += length + sharded = _get_tokens_on_this_cp_rank(tokens, cp_rank, cp_size, seq_dim=1) + assert reconstructed == sharded.squeeze(0).tolist() + + +def test_cp_load_balanced_segments_validate_partition(): + with pytest.raises(ValueError, match=r"divisible by 2\*cp_size"): + cp_load_balanced_segments(10, 0, 2) + with pytest.raises(ValueError, match="cp_rank must be"): + cp_load_balanced_segments(8, 2, 2) + + +def test_distributed_selected_logprobs_tp1_matches_dense_value_and_gradient(): + """TP1 exercises the real-vocab normalization and the custom backward.""" + torch.manual_seed(1234) + real_vocab_size = 5 + temperature = 0.7 + base = torch.randn(2, 3, 7, dtype=torch.float32) + # Padded LM-head columns would dominate if they leaked into logZ. + base[..., real_vocab_size:] += 100.0 + actual_logits = base.clone().requires_grad_(True) + reference_logits = base.clone().requires_grad_(True) + batch_indices = torch.tensor([0, 0, 1], dtype=torch.long) + seq_indices = torch.tensor([1, 1, 2], dtype=torch.long) + token_ids = torch.tensor([[0, 4], [4, 2], [1, 1]], dtype=torch.long) + weights = torch.tensor([[0.5, -0.25], [1.25, 0.75], [-0.5, 2.0]]) + + actual = distributed_selected_logprobs( + actual_logits, + batch_indices, + seq_indices, + token_ids, + vocab_start_index=0, + vocab_end_index=actual_logits.shape[-1], + real_vocab_size=real_vocab_size, + temperature=temperature, + tp_group=None, + row_chunk_size=1, + ) + dense_rows = reference_logits[batch_indices, seq_indices, :real_vocab_size] + expected = torch.gather( + F.log_softmax(dense_rows / temperature, dim=-1), dim=-1, index=token_ids + ) + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-6) + + (actual * weights).sum().backward() + (expected * weights).sum().backward() + torch.testing.assert_close( + actual_logits.grad, reference_logits.grad, rtol=2e-6, atol=2e-6 + ) + assert torch.equal( + actual_logits.grad[..., real_vocab_size:], + torch.zeros_like(actual_logits.grad[..., real_vocab_size:]), + ) + + +def test_distributed_selected_logprobs_marks_out_of_vocab_ids_minus_inf(): + logits = torch.randn(1, 2, 6) + actual = distributed_selected_logprobs( + logits, + torch.tensor([0, 0]), + torch.tensor([0, 1]), + torch.tensor([[-1, 2], [5, 4]]), + vocab_start_index=0, + vocab_end_index=6, + real_vocab_size=5, + temperature=1.0, + ) + assert actual[0, 0].item() == float("-inf") + assert actual[1, 0].item() == float("-inf") + assert torch.isfinite(actual[0, 1]) and torch.isfinite(actual[1, 1]) + + +def test_distributed_vocab_topk_logz_force_tp1_exports_natural_and_sidecars(): + """Natural K is unchanged while every forced label gets an exact sidecar.""" + logits = torch.tensor( + [ + [ + [9.0, 8.0, 7.0, 6.0, 5.0, 100.0, 90.0], + [9.0, 8.0, 7.0, 6.0, 5.0, 110.0, 80.0], + [1.0, 5.0, 4.0, 3.0, 2.0, 120.0, 70.0], + ] + ] + ) + force = torch.tensor([[[1, 4], [4, -1], [-1, -1]]], dtype=torch.long) + temperature = 2.0 + + ( + values, + indices, + log_z, + natural_tail_indices, + forced_values, + forced_indices, + forced_was_natural, + ) = distributed_vocab_topk_logz_force( + logits, + force, + 2, + None, + vocab_start_index=0, + vocab_end_index=logits.shape[-1], + real_vocab_size=5, + temperature=temperature, + chunk_size=2, + ) + + assert indices.tolist() == [[[0, 1], [0, 1], [1, 2]]] + torch.testing.assert_close( + values, + torch.tensor([[[9.0, 8.0], [9.0, 8.0], [5.0, 4.0]]]), + rtol=0, + atol=0, + ) + # Values stay raw; only the exact denominator is temperature scaled. + torch.testing.assert_close( + log_z, + torch.logsumexp(logits[..., :5] / temperature, dim=-1), + rtol=1e-6, + atol=1e-6, + ) + assert natural_tail_indices.tolist() == [[[1], [1], [2]]] + assert forced_indices.tolist() == force.tolist() + assert forced_values.tolist() == [[[8.0, 5.0], [5.0, 0.0], [0.0, 0.0]]] + assert forced_was_natural.tolist() == [ + [[True, False], [False, False], [False, False]] + ] + assert bool((indices < 5).all()) + + +def test_distributed_vocab_topk_logz_force_supports_k_one(): + logits = torch.tensor([[[4.0, 3.0, 2.0, 1.0, 0.0, 999.0]]]) + ( + values, + indices, + log_z, + natural_tail_indices, + forced_values, + forced_indices, + forced_was_natural, + ) = distributed_vocab_topk_logz_force( + logits, + torch.tensor([[[4, -1]]], dtype=torch.long), + 1, + None, + vocab_start_index=0, + vocab_end_index=6, + real_vocab_size=5, + temperature=0.5, + ) + + assert indices.tolist() == [[[0]]] + assert values.tolist() == [[[4.0]]] + assert natural_tail_indices.tolist() == [[[0]]] + assert forced_indices.tolist() == [[[4, -1]]] + assert forced_values.tolist() == [[[0.0, 0.0]]] + assert forced_was_natural.tolist() == [[[False, False]]] + torch.testing.assert_close( + log_z, + torch.logsumexp(logits[..., :5] / 0.5, dim=-1), + rtol=1e-6, + atol=1e-6, + ) + + +def test_distributed_vocab_topk_logz_force_membership_can_exceed_transport_k(): + logits = torch.tensor([[[5.0, 4.0, 3.0, 2.0, 1.0]]]) + outputs = distributed_vocab_topk_logz_force( + logits, + torch.tensor([[[2, 4]]], dtype=torch.long), + 2, + None, + vocab_start_index=0, + vocab_end_index=5, + real_vocab_size=5, + temperature=1.0, + noise_filter_k=3, + ) + + assert outputs[-1].tolist() == [[[True, False]]] + + +def test_distributed_vocab_topk_logz_force_canonicalizes_cutoff_ties(): + logits = torch.tensor([[[5.0, 4.0, 4.0, 1.0]]]) + outputs = distributed_vocab_topk_logz_force( + logits, + torch.tensor([[[3, -1]]], dtype=torch.long), + 2, + None, + vocab_start_index=0, + vocab_end_index=4, + real_vocab_size=4, + temperature=1.0, + ) + + # Equal cutoff logits keep the lowest global token id, independent of the + # candidate order a TP partition presents to torch.topk. + assert outputs[1].tolist() == [[[0, 1]]] + assert outputs[3].tolist() == [[[1]]] + + +def test_distributed_vocab_topk_logz_force_rejects_nan_temperature(): + with pytest.raises(ValueError, match="finite and > 0"): + distributed_vocab_topk_logz_force( + torch.tensor([[[5.0, 4.0, 3.0]]]), + torch.tensor([[[2, -1]]]), + 2, + None, + vocab_start_index=0, + vocab_end_index=3, + real_vocab_size=3, + temperature=float("nan"), + ) + + +@pytest.mark.parametrize( + "force,match", + [ + (torch.tensor([[[5, -1]]]), "real-vocabulary ids"), + (torch.tensor([[[2, 2]]]), "duplicate valid labels"), + ], +) +def test_distributed_vocab_topk_logz_force_rejects_invalid_sidecars(force, match): + with pytest.raises(ValueError, match=match): + distributed_vocab_topk_logz_force( + torch.tensor([[[5.0, 4.0, 3.0, 2.0, 1.0]]]), + force, + 2, + None, + vocab_start_index=0, + vocab_end_index=5, + real_vocab_size=5, + temperature=1.0, + ) + + +def _run_sparse_primitives_tp(rank: int, tp_size: int, init_file: str) -> None: + """Selected log-probs and sparse top-k on a gloo TP group match a dense view.""" + torch.set_num_threads(1) + dist.init_process_group( + backend="gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=tp_size, + ) + try: + tp_group = dist.new_group(ranks=list(range(tp_size))) + batch_size, seq_len = 2, 5 + padded_vocab_size, real_vocab_size = 8, 7 + local_vocab_size = padded_vocab_size // tp_size + vocab_start_index = rank * local_vocab_size + vocab_end_index = vocab_start_index + local_vocab_size + temperature = 0.7 + + torch.manual_seed(20260805) + full_logits = torch.randn(batch_size, seq_len, padded_vocab_size) + # Column 7 is TP-divisibility padding. It would dominate both top-k + # and logZ if the real-vocabulary boundary leaked. + full_logits[..., real_vocab_size] += 100.0 + local_logits = ( + full_logits[..., vocab_start_index:vocab_end_index] + .detach() + .clone() + .requires_grad_(True) + ) + + batch_indices = torch.tensor([0, 0, 1, 1]) + seq_indices = torch.tensor([1, 1, 2, 4]) + token_ids = torch.tensor([[0, 6], [5, 2], [3, 3], [6, 1]]) + weights = torch.tensor([[0.5, -0.25], [1.25, 0.75], [-0.5, 2.0], [0.3, -1.1]]) + selected = distributed_selected_logprobs( + local_logits, + batch_indices, + seq_indices, + token_ids, + vocab_start_index=vocab_start_index, + vocab_end_index=vocab_end_index, + real_vocab_size=real_vocab_size, + temperature=temperature, + tp_group=tp_group, + row_chunk_size=2, + ) + dense_for_grad = full_logits.detach().clone().requires_grad_(True) + dense_rows = dense_for_grad[batch_indices, seq_indices, :real_vocab_size] + expected_selected = torch.gather( + torch.log_softmax(dense_rows / temperature, dim=-1), + dim=-1, + index=token_ids, + ) + torch.testing.assert_close(selected, expected_selected, rtol=2e-6, atol=2e-6) + (selected * weights).sum().backward() + (expected_selected * weights).sum().backward() + + gathered_grads = [torch.empty_like(local_logits.grad) for _ in range(tp_size)] + dist.all_gather(gathered_grads, local_logits.grad, group=tp_group) + reconstructed_grad = torch.cat(gathered_grads, dim=-1) + torch.testing.assert_close( + reconstructed_grad, dense_for_grad.grad, rtol=3e-6, atol=3e-6 + ) + assert torch.equal( + reconstructed_grad[..., real_vocab_size], + torch.zeros_like(reconstructed_grad[..., real_vocab_size]), + ) + + # Rows request a natural top-3 label, a deliberately absent + # (minimum-logit) label, or no label. A second distinct sidecar entry + # exercises forced logits owned by different TP ranks. + real_logits = full_logits[..., :real_vocab_size] + natural_top3 = torch.topk(real_logits, 3, dim=-1).indices + primary_force_ids = natural_top3[..., 0].clone() + absent_rows = torch.zeros_like(primary_force_ids, dtype=torch.bool) + absent_rows[0, 1] = True + absent_rows[1, 3] = True + primary_force_ids[absent_rows] = real_logits.argmin(dim=-1)[absent_rows] + secondary_force_ids = (primary_force_ids + 4) % real_vocab_size + # One row spans the two TP vocabulary shards. + primary_force_ids[0, 0] = 0 + secondary_force_ids[0, 0] = 6 + primary_force_ids[0, 4] = -1 + secondary_force_ids[0, 4] = -1 + force_ids = torch.stack([primary_force_ids, secondary_force_ids], dim=-1) + + ( + topk_values, + topk_ids, + log_z, + natural_tail_ids, + forced_values, + transported_force_ids, + forced_was_natural, + ) = distributed_vocab_topk_logz_force( + local_logits.detach(), + force_ids, + 2, + tp_group, + vocab_start_index=vocab_start_index, + vocab_end_index=vocab_end_index, + real_vocab_size=real_vocab_size, + temperature=temperature, + noise_filter_k=3, + chunk_size=3, + ) + expected_values, expected_ids = torch.topk(real_logits, 2, dim=-1) + expected_tail_ids = expected_ids[..., -1:].clone() + force_valid = force_ids >= 0 + expected_membership = ( + natural_top3.unsqueeze(-2) == force_ids.unsqueeze(-1) + ).any(dim=-1) & force_valid + expected_forced_values = torch.where( + force_valid, + torch.gather(real_logits, -1, force_ids.clamp_min(0)), + torch.zeros(force_ids.shape), + ) + expected_ids, order = torch.sort(expected_ids, dim=-1) + expected_values = torch.gather(expected_values, -1, order) + + torch.testing.assert_close(topk_values, expected_values, rtol=0, atol=0) + assert torch.equal(topk_ids.long(), expected_ids) + assert torch.equal(natural_tail_ids.long(), expected_tail_ids) + torch.testing.assert_close( + forced_values, expected_forced_values, rtol=0, atol=0 + ) + assert torch.equal(transported_force_ids.long(), force_ids) + torch.testing.assert_close( + log_z, + torch.logsumexp(real_logits / temperature, dim=-1), + rtol=2e-6, + atol=2e-6, + ) + assert torch.equal(forced_was_natural, expected_membership) + assert bool((topk_ids < real_vocab_size).all()) + + # Exact cutoff ties select the same lowest global token ids at every + # TP size, although each shard presents its own candidates. + tie_full_logits = torch.tensor([[[9.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 999.0]]]) + tie_outputs = distributed_vocab_topk_logz_force( + tie_full_logits[..., vocab_start_index:vocab_end_index], + torch.tensor([[[6, -1]]], dtype=torch.long), + 3, + tp_group, + vocab_start_index=vocab_start_index, + vocab_end_index=vocab_end_index, + real_vocab_size=real_vocab_size, + temperature=1.0, + chunk_size=1, + ) + assert tie_outputs[1].long().tolist() == [[[0, 1, 2]]] + assert tie_outputs[3].long().tolist() == [[[2]]] + torch.testing.assert_close( + tie_outputs[0], torch.tensor([[[9.0, 8.0, 8.0]]]), rtol=0, atol=0 + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not dist.is_available() or not dist.is_gloo_available(), + reason="gloo backend unavailable", +) +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_sparse_primitives_match_dense_on_gloo_tp(tmp_path, tp_size): + mp.spawn( + _run_sparse_primitives_tp, + args=(tp_size, os.fspath(tmp_path / "gloo_init")), + nprocs=tp_size, + join=True, + ) diff --git a/tests/unit/distributed/test_student_loss_primitives.py b/tests/unit/distributed/test_student_loss_primitives.py new file mode 100644 index 00000000000..10cb7903890 --- /dev/null +++ b/tests/unit/distributed/test_student_loss_primitives.py @@ -0,0 +1,56 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from nemo_rl.distributed.model_utils import chunked_vocab_parallel_logsumexp + + +def test_chunked_vocab_logsumexp_tp1_matches_across_chunk_boundaries(): + """TP1 keeps exact forward/gradient math for rows around the fixed cap.""" + torch.manual_seed(7) + logits = torch.randn(2, 513, 17, dtype=torch.float64, requires_grad=True) + reference_logits = logits.detach().clone().requires_grad_(True) + weights = torch.randn(2, 513, dtype=torch.float64) + + actual = chunked_vocab_parallel_logsumexp( + logits, + 1.7, + chunk_size=256, + ) + expected = torch.logsumexp(reference_logits / 1.7, dim=-1) + assert torch.equal(actual, expected) + + (actual * weights).sum().backward() + (expected * weights).sum().backward() + assert torch.equal(logits.grad, reference_logits.grad) + + +@pytest.mark.parametrize( + ("temperature", "chunk_size", "match"), + [(0.0, 256, "temperature"), (1.0, 0, "chunk_size")], +) +def test_chunked_vocab_logsumexp_rejects_invalid_bounds( + temperature: float, + chunk_size: int, + match: str, +): + logits = torch.zeros(1, 2, 3) + with pytest.raises(ValueError, match=match): + chunked_vocab_parallel_logsumexp( + logits, + temperature, + chunk_size=chunk_size, + ) diff --git a/tests/unit/models/automodel/test_automodel_data.py b/tests/unit/models/automodel/test_automodel_data.py index 562edece2c8..f0acf3d9ff4 100644 --- a/tests/unit/models/automodel/test_automodel_data.py +++ b/tests/unit/models/automodel/test_automodel_data.py @@ -19,6 +19,11 @@ from nemo_rl.algorithms.loss.interfaces import LossType from nemo_rl.data.multimodal_utils import PackedTensor +from nemo_rl.data.packing import ( + LockstepPackingItem, + SidePackingSpec, + build_lockstep_packing_plan, +) from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.automodel.data import ( ProcessedInputs, @@ -287,6 +292,131 @@ def side_effect(tensor, *args, **kwargs): @pytest.mark.automodel class TestProcessMicrobatch: + @staticmethod + def _planned_microbatch(*, fixed_tail: bool = True): + plan = build_lockstep_packing_plan( + batch_uid=17, + items=( + LockstepPackingItem(sample_id="a", batch_item_id=10), + LockstepPackingItem(sample_id="b", batch_item_id=11), + ), + sides=( + SidePackingSpec( + side_id="student", + capacity=12, + raw_lengths=(3, 2), + effective_lengths=(4, 4), + physical_size_fn=(lambda _lengths: 12) if fixed_tail else None, + ), + ), + data_parallel_size=1, + ) + mb = BatchedDataDict( + { + "input_ids": torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]]), + "input_lengths": torch.tensor([3, 2]), + "sample_mask": torch.ones(2), + "batch_item_id": torch.tensor([10, 11]), + "kd_token_mask": torch.tensor( + [[0, 1, 1, 0], [0, 1, 0, 0]], dtype=torch.bool + ), + "student_semantic_regions": [ + ((10, "assistant", "eot", 1, 3),), + ((11, "assistant", "eot", 1, 2),), + ], + } + ) + mb.lockstep_packing_plan = plan + mb.lockstep_side_id = "student" + mb.lockstep_bin_indices = (0,) + return mb + + @patch("torch.Tensor.cuda", lambda self: self) + def test_plan_driven_packing_materializes_exact_raw_and_padded_geometry( + self, mock_tokenizer + ): + mb = self._planned_microbatch() + semantic_regions = mb["student_semantic_regions"] + + result = process_microbatch( + mb=mb, + tokenizer=mock_tokenizer, + enable_seq_packing=True, + cfg={ + "automodel_cfg": { + "sequence_parallel": False, + "context_parallel_size": 1, + "tensor_parallel_size": 1, + }, + "sequence_packing": {"train_mb_tokens": 999}, + }, + ) + + assert result.input_ids.tolist() == [[1, 2, 3, 2, 4, 5, 2, 2, 2, 2, 2, 2]] + assert result.position_ids.tolist() == [[0, 1, 2, 0, 0, 1, 0, 0, 0, 0, 0, 0]] + assert result.flash_attn_kwargs.cu_seqlens_q.tolist() == [0, 3, 5] + assert result.flash_attn_kwargs.cu_seqlens_q_padded.tolist() == [0, 4, 12] + assert result.flash_attn_kwargs.cu_seq_lens_q.tolist() == [0, 4, 12] + assert result.seq_len == 12 + # Logical metadata is retained for MBS1 loss calls and never forwarded + # through the model-input materializer. + assert mb["student_semantic_regions"] == semantic_regions + assert mb["kd_token_mask"].shape == (2, 4) + + @patch("torch.Tensor.cuda", lambda self: self) + def test_plan_driven_packing_rejects_order_and_length_drift(self, mock_tokenizer): + cfg = { + "automodel_cfg": { + "sequence_parallel": False, + "context_parallel_size": 1, + "tensor_parallel_size": 1, + }, + "sequence_packing": {"train_mb_tokens": 12}, + } + wrong_order = self._planned_microbatch() + wrong_order["batch_item_id"] = torch.tensor([11, 10]) + with pytest.raises(ValueError, match="batch_item_id order"): + process_microbatch(wrong_order, mock_tokenizer, True, cfg) + + wrong_length = self._planned_microbatch() + wrong_length["input_lengths"] = torch.tensor([2, 2]) + with pytest.raises(ValueError, match="input lengths drifted"): + process_microbatch(wrong_length, mock_tokenizer, True, cfg) + + @patch("torch.Tensor.cuda", lambda self: self) + def test_plan_driven_tp_training_requires_fixed_physical_tail(self, mock_tokenizer): + mb = self._planned_microbatch(fixed_tail=False) + with pytest.raises(ValueError, match="requires every student bin"): + process_microbatch( + mb=mb, + tokenizer=mock_tokenizer, + enable_seq_packing=True, + cfg={ + "automodel_cfg": { + "sequence_parallel": False, + "context_parallel_size": 1, + "tensor_parallel_size": 2, + }, + "sequence_packing": {"train_mb_tokens": 12}, + }, + ) + + fixed_tail_mb = self._planned_microbatch() + result = process_microbatch( + mb=fixed_tail_mb, + tokenizer=mock_tokenizer, + enable_seq_packing=True, + cfg={ + "automodel_cfg": { + "sequence_parallel": False, + "context_parallel_size": 1, + "tensor_parallel_size": 2, + }, + "sequence_packing": {"train_mb_tokens": 12}, + }, + ) + assert result.seq_len == 12 + def test_regular_batching(self, mock_tokenizer): # Create test microbatch mb = BatchedDataDict( @@ -952,6 +1082,51 @@ def side_effect(tensor, *args, **kwargs): # Verify all_reduce was called mock_all_reduce.assert_called_once() + @patch("torch.Tensor.cuda", lambda self: self) + @patch("nemo_rl.models.automodel.data.torch.distributed.all_reduce") + def test_global_chunk_counts_by_teacher( + self, mock_all_reduce, mock_loss_fn, mock_dp_mesh + ): + sample_mask = torch.tensor([1, 0, 1], dtype=torch.bool) + batch = BatchedDataDict( + { + "input_ids": torch.zeros(3, 4, dtype=torch.long), + "sample_mask": sample_mask, + "kd_token_mask": torch.tensor( + [[0, 1, 0, 1], [1, 1, 1, 1], [0, 1, 1, 0]], + dtype=torch.bool, + ), + "alignment_0_pair_valid": torch.tensor( + [[1, 1, 0], [1, 1, 1], [1, 0, 0]], dtype=torch.bool + ), + "alignment_2_pair_valid": torch.tensor( + [[1, 1, 1, 0], [1, 1, 1, 1], [1, 1, 0, 0]], + dtype=torch.bool, + ), + } + ) + data = MagicMock() + data.get_batch.return_value = batch + + def double_counts(values, *args, **kwargs): + torch.testing.assert_close(values.cpu(), torch.tensor([2, 8, 4, 3, 5])) + values.mul_(2) + + mock_all_reduce.side_effect = double_counts + + result = process_global_batch( + data=data, + loss_fn=mock_loss_fn, + dp_group=mock_dp_mesh.get_group(), + batch_idx=0, + batch_size=3, + ) + + assert { + i: count.item() for i, count in result["global_valid_chunks_by_idx"].items() + } == {0: 6, 2: 10} + assert result["global_valid_kd_toks"].item() == 8 + @patch("nemo_rl.models.automodel.data.torch.distributed.all_reduce") def test_with_token_mask(self, mock_all_reduce, mock_loss_fn, mock_dp_mesh): # Create test data diff --git a/tests/unit/models/automodel/test_automodel_setup.py b/tests/unit/models/automodel/test_automodel_setup.py index 93a827b1878..53412b3601c 100644 --- a/tests/unit/models/automodel/test_automodel_setup.py +++ b/tests/unit/models/automodel/test_automodel_setup.py @@ -36,6 +36,7 @@ ModelAndOptimizerState, RuntimeConfig, _maybe_set_force_hf, + _precision_to_dtype, get_tokenizer, setup_distributed, setup_model_and_optimizer, @@ -603,6 +604,115 @@ def test_missing_hf_config_overrides_becomes_empty_dict( assert result.hf_config_overrides == {} + def test_precision_to_dtype_supported(self): + """_precision_to_dtype maps each supported name to the right torch.dtype.""" + assert _precision_to_dtype("float32", "precision") == torch.float32 + assert _precision_to_dtype("bfloat16", "precision") == torch.bfloat16 + assert _precision_to_dtype("float16", "precision") == torch.float16 + + def test_precision_to_dtype_invalid(self): + """_precision_to_dtype raises ValueError on unknown name.""" + with pytest.raises(ValueError, match="Unknown automodel_cfg.load_precision"): + _precision_to_dtype("fp4", "automodel_cfg.load_precision") + + @patch("nemo_rl.models.automodel.setup.AutoConfig") + @patch("nemo_rl.models.automodel.setup.resolve_model_class") + @patch("nemo_rl.models.automodel.setup.configure_dynamo_cache") + def test_load_precision_default_preserves_automatic_dtype( + self, + mock_dynamo, + mock_resolve_class, + mock_autoconfig_class, + mock_config, + mock_autoconfig, + ): + """Unset load_precision preserves optimizer-aware dtype selection.""" + mock_autoconfig_class.from_pretrained.return_value = mock_autoconfig + mock_resolve_class.return_value = Mock + assert "load_precision" not in mock_config["automodel_cfg"] + + result = validate_and_prepare_config( + config=mock_config, + processor=None, + rank=0, + init_optimizer=True, + ) + assert result.model_load_dtype is None + + @patch("nemo_rl.models.automodel.setup.AutoConfig") + @patch("nemo_rl.models.automodel.setup.resolve_model_class") + @patch("nemo_rl.models.automodel.setup.configure_dynamo_cache") + def test_load_precision_bfloat16_reference_passes( + self, + mock_dynamo, + mock_resolve_class, + mock_autoconfig_class, + mock_config, + mock_autoconfig, + ): + """Reference/teacher workers (init_optimizer=False) can load in bfloat16.""" + mock_autoconfig_class.from_pretrained.return_value = mock_autoconfig + mock_resolve_class.return_value = Mock + mock_config["automodel_cfg"]["load_precision"] = "bfloat16" + + result = validate_and_prepare_config( + config=mock_config, + processor=None, + rank=0, + init_optimizer=False, + ) + assert result.model_load_dtype == torch.bfloat16 + + @patch("nemo_rl.models.automodel.setup.AutoConfig") + @patch("nemo_rl.models.automodel.setup.resolve_model_class") + @patch("nemo_rl.models.automodel.setup.configure_dynamo_cache") + def test_load_precision_bfloat16_trainable_raises( + self, + mock_dynamo, + mock_resolve_class, + mock_autoconfig_class, + mock_config, + ): + """Trainable workers (init_optimizer=True) must use float32; bfloat16 raises.""" + mock_config["automodel_cfg"]["load_precision"] = "bfloat16" + + with pytest.raises( + ValueError, + match="load_precision must be float32 for trainable workers", + ): + validate_and_prepare_config( + config=mock_config, + processor=None, + rank=0, + init_optimizer=True, + ) + + @patch("nemo_rl.models.automodel.setup.AutoConfig") + @patch("nemo_rl.models.automodel.setup.resolve_model_class") + @patch("nemo_rl.models.automodel.setup.configure_dynamo_cache") + def test_load_precision_passed_to_autoconfig_from_pretrained( + self, + mock_dynamo, + mock_resolve_class, + mock_autoconfig_class, + mock_config, + mock_autoconfig, + ): + """model_load_dtype is forwarded to AutoConfig.from_pretrained as torch_dtype.""" + mock_autoconfig_class.from_pretrained.return_value = mock_autoconfig + mock_resolve_class.return_value = Mock + mock_config["automodel_cfg"]["load_precision"] = "bfloat16" + + validate_and_prepare_config( + config=mock_config, + processor=None, + rank=0, + init_optimizer=False, + ) + assert mock_autoconfig_class.from_pretrained.called + kwargs = mock_autoconfig_class.from_pretrained.call_args.kwargs + assert kwargs["torch_dtype"] == torch.bfloat16 + @pytest.mark.automodel class TestSetupReferenceModelState: diff --git a/tests/unit/models/automodel/test_automodel_train.py b/tests/unit/models/automodel/test_automodel_train.py index 46b8623836f..7aa798a90f3 100644 --- a/tests/unit/models/automodel/test_automodel_train.py +++ b/tests/unit/models/automodel/test_automodel_train.py @@ -35,10 +35,12 @@ make_processed_microbatch_iterator, ) from nemo_rl.models.automodel.train import ( + FullLogitsPostProcessor, LogprobsPostProcessor, LossPostProcessor, PreparedModelForward, ScorePostProcessor, + SparseLogitsPostProcessor, TopkLogitsPostProcessor, apply_temperature_scaling, automodel_forward_backward, @@ -648,6 +650,49 @@ def test_loss_with_sequence_packing( # Verify the wrapper was called instead of raw loss_fn mock_wrapper_instance.assert_called_once() + @patch("nemo_rl.models.automodel.train.SequencePackingLossWrapper") + @patch("nemo_rl.models.automodel.train.CSDSequencePackingLossWrapper") + def test_csd_packing_selects_axis_aware_wrapper_only( + self, + mock_csd_wrapper_class, + mock_generic_wrapper_class, + base_cfg, + mock_cp_mesh, + processed_inputs_with_flash, + ): + from nemo_rl.algorithms.loss.csd import CSDLossFn + + loss_fn = object.__new__(CSDLossFn) + wrapper = MagicMock(return_value=(torch.tensor(0.5), {})) + mock_csd_wrapper_class.return_value = wrapper + processor = LossPostProcessor( + loss_fn=loss_fn, + cfg=base_cfg, + cp_mesh=mock_cp_mesh, + cp_size=1, + dp_size=1, + enable_seq_packing=True, + ) + data = BatchedDataDict( + { + "input_ids": torch.zeros(1, 8, dtype=torch.long), + "sample_mask": torch.ones(1), + } + ) + + processor( + logits=torch.randn(1, 8, 4), + data_dict=data, + processed_inputs=processed_inputs_with_flash, + global_valid_seqs=torch.tensor(1.0), + global_valid_toks=torch.tensor(1.0), + cp_sharder=None, + ) + + mock_csd_wrapper_class.assert_called_once() + mock_generic_wrapper_class.assert_not_called() + wrapper.assert_called_once() + def test_loss_processor_initialization( self, base_cfg, @@ -825,6 +870,75 @@ def test_basic_topk(self, base_cfg, mock_device_mesh, mock_cp_mesh, mock_tp_mesh assert idx.shape == (batch_size, seq_len, k) +# ===================== +# Test SparseLogitsPostProcessor +# ===================== +@pytest.mark.automodel +class TestSparseLogitsPostProcessor: + def test_row_topk_logz_padding_and_forced_label(self, base_cfg, monkeypatch): + """Sparse export keeps exact logZ and ignores padded vocab columns.""" + group = object() + mesh = MagicMock() + mesh.get_group.return_value = group + + monkeypatch.setattr(torch.distributed, "get_rank", lambda _group: 0) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda _group: 1) + monkeypatch.setattr( + torch.distributed, + "all_reduce", + lambda tensor, op=None, group=None: tensor, + ) + + def _all_gather(outputs, tensor, group=None): + assert len(outputs) == 1 + outputs[0].copy_(tensor) + + monkeypatch.setattr(torch.distributed, "all_gather", _all_gather) + + processor = SparseLogitsPostProcessor( + cfg=base_cfg, + tp_mesh=mesh, + cp_mesh=mesh, + cp_size=1, + k=3, + temperature=2.0, + vocab_size=5, + gt_filter_topk=2, + ) + # Column 5 is padded lm-head vocabulary. Its enormous value must not + # affect top-k support or the exact full-real-vocabulary normalizer. + logits = torch.tensor( + [[[5.0, 4.0, 3.0, 2.0, -7.0, 1000.0], [0.0, 1.0, 2.0, 3.0, 4.0, 900.0]]] + ) + data = BatchedDataDict( + {"force_include_token_ids": torch.tensor([[4, 4]], dtype=torch.long)} + ) + + values, token_ids, log_z, gt_in_topk = processor( + logits=logits, + data_dict=data, + processed_inputs=MagicMock(), + original_batch_size=1, + original_seq_len=2, + cp_sharder=None, + ) + + assert torch.equal( + token_ids, + torch.tensor([[[0, 1, 4], [2, 3, 4]]], dtype=torch.int32), + ) + torch.testing.assert_close( + values, + torch.tensor([[[5.0, 4.0, -7.0], [2.0, 3.0, 4.0]]]), + ) + torch.testing.assert_close( + log_z, + torch.logsumexp(logits[..., :5] / 2.0, dim=-1), + ) + assert gt_in_topk is not None + assert torch.equal(gt_in_topk, torch.tensor([[False, True]])) + + # ===================== # Test ProcessedMicrobatch # ===================== @@ -1255,7 +1369,6 @@ def test_forward_backward_single_microbatch( def test_forward_backward_multiple_microbatches( self, mock_model, - mock_loss_fn, base_cfg, mock_device_mesh, mock_cp_mesh, @@ -1296,9 +1409,16 @@ def test_forward_backward_multiple_microbatches( ) ) + from nemo_rl.algorithms.loss.csd import CSDLossFn + + # Only CSD losses receive the global alignment-chunk counts. + csd_loss_fn = MagicMock(spec=CSDLossFn) + csd_loss_fn.return_value = (torch.tensor(0.5), {"loss": 0.5}) + csd_loss_fn.input_type = LossInputType.LOGIT + # Create loss post-processor loss_post_processor = LossPostProcessor( - loss_fn=mock_loss_fn, + loss_fn=csd_loss_fn, cfg=base_cfg, cp_mesh=mock_cp_mesh, cp_size=1, @@ -1306,6 +1426,7 @@ def test_forward_backward_multiple_microbatches( ) # Call automodel_forward_backward in forward_only mode + global_valid_chunks_by_idx = {0: torch.tensor(37)} results = automodel_forward_backward( model=mock_model, data_iterator=iter(processed_mbs), @@ -1316,6 +1437,7 @@ def test_forward_backward_multiple_microbatches( forward_only=True, global_valid_seqs=torch.tensor(batch_size * num_microbatches), global_valid_toks=torch.tensor(batch_size * seq_len * num_microbatches), + global_valid_chunks_by_idx=global_valid_chunks_by_idx, ) # Verify results @@ -1325,7 +1447,11 @@ def test_forward_backward_multiple_microbatches( assert mock_model.call_count == num_microbatches # Verify loss function was called num_microbatches times - assert mock_loss_fn.call_count == num_microbatches + assert csd_loss_fn.call_count == num_microbatches + assert all( + call.kwargs["global_valid_chunks_by_idx"] is global_valid_chunks_by_idx + for call in csd_loss_fn.call_args_list + ) @patch("nemo_rl.models.automodel.train.prepare_model_forward") def test_forward_backward_enters_prepared_and_autocast_contexts( @@ -2094,6 +2220,57 @@ class MockFlashAttnKwargs: torch.testing.assert_close(idx, expected_idx) +@pytest.mark.automodel +class TestFullLogitsPostProcessorSeqPacking: + def test_restores_one_dense_row_per_logical_sample(self): + processor = FullLogitsPostProcessor( + cfg={"sequence_packing": {"enabled": True}}, + cp_mesh=None, + cp_size=1, + enable_seq_packing=True, + ) + processed_inputs = MagicMock() + processed_inputs.flash_attn_kwargs.cu_seqlens_q = torch.tensor( + [0, 3, 5], dtype=torch.int32 + ) + processed_inputs.flash_attn_kwargs.cu_seqlens_q_padded = torch.tensor( + [0, 4, 6], dtype=torch.int32 + ) + data = BatchedDataDict( + { + "input_ids": torch.zeros(2, 6, dtype=torch.long), + "input_lengths": torch.tensor([3, 2]), + } + ) + packed_logits = torch.tensor( + [ + [ + [1.0, 101.0], + [2.0, 102.0], + [3.0, 103.0], + [99.0, 199.0], + [10.0, 110.0], + [11.0, 111.0], + ] + ] + ) + + restored = processor( + logits=packed_logits, + data_dict=data, + processed_inputs=processed_inputs, + original_batch_size=2, + original_seq_len=6, + cp_sharder=None, + ) + + assert restored.shape == (2, 6, 2) + assert torch.equal(restored[0, :3], packed_logits[0, :3].float()) + assert torch.equal(restored[1, :2], packed_logits[0, 4:6].float()) + assert torch.count_nonzero(restored[0, 3:]) == 0 + assert torch.count_nonzero(restored[1, 2:]) == 0 + + # ===================== # Test automodel_forward_backward with backward pass # ===================== diff --git a/tests/unit/models/megatron/test_megatron_data.py b/tests/unit/models/megatron/test_megatron_data.py index 0b64214167a..6ac68ba1c53 100644 --- a/tests/unit/models/megatron/test_megatron_data.py +++ b/tests/unit/models/megatron/test_megatron_data.py @@ -29,6 +29,11 @@ import ray import torch +from nemo_rl.data.packing import ( + LockstepPackingItem, + SidePackingSpec, + build_lockstep_packing_plan, +) from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.named_sharding import NamedSharding from nemo_rl.distributed.ray_actor_environment_registry import ( @@ -469,6 +474,10 @@ def test_process_microbatch_with_packing( data_dict.__contains__ = MagicMock( side_effect=lambda k: k in {"input_ids", "input_lengths"} ) + data_dict.lockstep_packing_plan = None + data_dict.lockstep_side_id = None + data_dict.lockstep_bin_indices = None + data_dict.get = MagicMock(return_value=None) result = process_microbatch( data_dict, @@ -1043,6 +1052,102 @@ def test_caller_packing_matches_mbridge_thd_contract(self): assert torch.equal(cu_seqlens, mbridge_batch["cu_seqlens_q"]) assert torch.equal(cu_seqlens_padded, mbridge_batch["cu_seqlens_q_padded"]) + def test_process_microbatch_no_packing_context_parallel_shards_sequence_data( + self, + ): + """Unpacked CP uses the same zigzag shard for tokens, positions, and MTP.""" + from nemo_rl.models.megatron.data import process_microbatch + + input_ids = torch.arange(8).unsqueeze(0) + mtp_loss_mask = torch.arange(200, 208).unsqueeze(0) + + with ( + patch( + "nemo_rl.models.megatron.data.get_context_parallel_world_size", + return_value=2, + ), + patch( + "nemo_rl.models.megatron.data.get_context_parallel_rank", + return_value=1, + ), + ): + result = process_microbatch( + {"input_ids": input_ids, "mtp_loss_mask": mtp_loss_mask}, + pack_sequences=False, + straggler_timer=MagicMock(), + ) + + # CP rank 1 of 2 owns load-balanced chunks 1 and 2 of 4. + assert result.input_ids_cp_sharded.tolist() == [[2, 3, 4, 5]] + assert result.position_ids.tolist() == [[2, 3, 4, 5]] + assert result.mtp_loss_mask.tolist() == [[202, 203, 204, 205]] + assert result.attention_mask is None + + def test_process_microbatch_no_packing_context_parallel_shards_router_padding_mask( + self, + ): + """Unpacked CP shards the router padding mask like input_ids.""" + from nemo_rl.models.megatron.data import process_microbatch + + input_ids = torch.arange(8).unsqueeze(0) + + with ( + patch( + "nemo_rl.models.megatron.data.get_context_parallel_world_size", + return_value=2, + ), + patch( + "nemo_rl.models.megatron.data.get_context_parallel_rank", + return_value=0, + ), + ): + result = process_microbatch( + {"input_ids": input_ids, "input_lengths": torch.tensor([6])}, + pack_sequences=False, + straggler_timer=MagicMock(), + create_nonpacked_router_padding_mask=True, + ) + + # CP rank 0 of 2 owns load-balanced chunks 0 and 3 of 4; tokens 6-7 pad. + assert result.input_ids_cp_sharded.tolist() == [[0, 1, 6, 7]] + assert result.padding_mask.shape == result.input_ids_cp_sharded.shape + assert result.padding_mask.tolist() == [[False, False, True, True]] + + @pytest.mark.parametrize("cp_size", [1, 2]) + def test_process_microbatch_no_packing_builds_dense_mask_only_without_cp( + self, cp_size + ): + """Unpacked CP>1 never materializes the dense [1, 1, S, S] causal mask.""" + from nemo_rl.models.megatron import data as megatron_data + + with ( + patch( + "nemo_rl.models.megatron.data.get_context_parallel_world_size", + return_value=cp_size, + ), + patch( + "nemo_rl.models.megatron.data.get_context_parallel_rank", + return_value=0, + ), + patch( + "nemo_rl.models.megatron.data.get_ltor_masks_and_position_ids", + wraps=megatron_data.get_ltor_masks_and_position_ids, + ) as get_masks, + ): + result = megatron_data.process_microbatch( + {"input_ids": torch.arange(8).unsqueeze(0)}, + pack_sequences=False, + straggler_timer=MagicMock(), + ) + + # Megatron defaults create_attention_mask to True. + create_mask = get_masks.call_args.kwargs.get("create_attention_mask", True) + assert create_mask is (cp_size == 1) + if cp_size == 1: + assert result.attention_mask.shape == (1, 1, 8, 8) + else: + assert result.attention_mask is None + @patch("nemo_rl.models.megatron.data.get_ltor_masks_and_position_ids") def test_process_microbatch_no_packing_mtp_loss_mask_absent(self, mock_get_masks): """mtp_loss_mask defaults to None when not provided.""" @@ -1428,6 +1533,9 @@ def test_process_microbatch_delegate_pack_to_model(self, mock_prepare, mock_pack data_dict.__contains__ = MagicMock( side_effect=lambda k: k in {"input_ids", "input_lengths"} ) + data_dict.lockstep_packing_plan = None + data_dict.lockstep_side_id = None + data_dict.lockstep_bin_indices = None result = process_microbatch( data_dict, @@ -1471,6 +1579,9 @@ def test_process_microbatch_delegate_pack_rejects_mtp_loss_mask(self): data_dict.__contains__ = MagicMock( side_effect=lambda k: k in {"input_ids", "input_lengths", "mtp_loss_mask"} ) + data_dict.lockstep_packing_plan = None + data_dict.lockstep_side_id = None + data_dict.lockstep_bin_indices = None with pytest.raises(AssertionError) as exc_info: process_microbatch( @@ -1773,6 +1884,53 @@ def test_process_global_batch_basic(self): # Verify all_reduce was called mock_all_reduce.assert_called_once() + @patch("torch.Tensor.cuda", lambda self: self) + def test_process_global_batch_counts_chunks_by_teacher(self): + """Chunk counts exclude masked samples and remain teacher-specific.""" + from nemo_rl.models.megatron.data import process_global_batch + + sample_mask = torch.tensor([1, 0, 1], dtype=torch.bool) + batch = BatchedDataDict( + { + "input_ids": torch.zeros(3, 4, dtype=torch.long), + "sample_mask": sample_mask, + "kd_token_mask": torch.tensor( + [[0, 1, 0, 1], [1, 1, 1, 1], [0, 1, 1, 0]], + dtype=torch.bool, + ), + "alignment_0_pair_valid": torch.tensor( + [[1, 1, 0], [1, 1, 1], [1, 0, 0]], dtype=torch.bool + ), + "alignment_2_pair_valid": torch.tensor( + [[1, 1, 1, 0], [1, 1, 1, 1], [1, 1, 0, 0]], + dtype=torch.bool, + ), + } + ) + data = MagicMock() + data.get_batch.return_value = batch + + def double_counts(values, *args, **kwargs): + torch.testing.assert_close(values.cpu(), torch.tensor([2, 8, 4, 3, 5])) + values.mul_(2) + + with patch( + "torch.distributed.all_reduce", side_effect=double_counts + ) as mock_all_reduce: + result = process_global_batch( + data=data, + loss_fn=MagicMock(), + dp_group=MagicMock(), + batch_idx=0, + batch_size=3, + ) + + mock_all_reduce.assert_called_once() + assert { + i: count.item() for i, count in result["global_valid_chunks_by_idx"].items() + } == {0: 6, 2: 10} + assert result["global_valid_kd_toks"].item() == 8 + def test_process_global_batch_requires_sample_mask_in_data(self): """Test that process_global_batch requires sample_mask.""" from nemo_rl.models.megatron.data import process_global_batch @@ -2787,7 +2945,7 @@ def test_shard_routed_experts_for_cp_matches_input_ids_zigzag(cp_size): ( _all_input_ids, input_ids_cp_sharded, - _packed_seq_params, + packed_seq_params, cu_seqlens, cu_seqlens_padded, ) = _pack_sequences_for_megatron( @@ -2798,6 +2956,14 @@ def test_shard_routed_experts_for_cp_matches_input_ids_zigzag(cp_size): cp_size=cp_size, ) + # Raw boundaries drive logical sequence semantics (e.g. MTP and packed + # loss slicing); padded boundaries drive THD/CP physical placement. + assert not torch.equal(cu_seqlens, cu_seqlens_padded) + assert torch.equal(packed_seq_params.cu_seqlens_q, cu_seqlens) + assert torch.equal(packed_seq_params.cu_seqlens_kv, cu_seqlens) + assert torch.equal(packed_seq_params.cu_seqlens_q_padded, cu_seqlens_padded) + assert torch.equal(packed_seq_params.cu_seqlens_kv_padded, cu_seqlens_padded) + ( routed_packed, routed_cp_sharded, @@ -2837,6 +3003,121 @@ def test_shard_routed_experts_for_cp_matches_input_ids_zigzag(cp_size): assert torch.equal(routed_packed[0, pad_pos, layer], expected_route) +@pytest.mark.mcore +@pytest.mark.parametrize( + ("cp_size", "expected"), + [ + ( + 1, + [ + [ + False, + False, + False, + True, + True, + True, + True, + True, + True, + True, + True, + True, + ] + ], + ), + (2, [[False, True, True, True, True, True]]), + ], + ids=["cp1", "cp2"], +) +def test_packed_router_padding_mask_excludes_tail_and_masked_logical_row( + cp_size, expected +): + from nemo_rl.models.megatron.data import _make_packed_router_padding_mask + + padding_mask = _make_packed_router_padding_mask( + torch.tensor([3, 2]), + torch.tensor([0, 4, 12], dtype=torch.int32), + sample_mask=torch.tensor([1, 0]), + cp_rank=0, + cp_size=cp_size, + device=torch.device("cpu"), + ) + + assert padding_mask.tolist() == expected + + +@pytest.mark.mcore +def test_process_microbatch_validates_and_materializes_lockstep_geometry(): + from nemo_rl.models.megatron.data import process_microbatch + + plan = build_lockstep_packing_plan( + batch_uid=23, + items=( + LockstepPackingItem(sample_id="a", batch_item_id=10), + LockstepPackingItem(sample_id="b", batch_item_id=11), + ), + sides=( + SidePackingSpec( + side_id="student", + capacity=12, + raw_lengths=(3, 2), + effective_lengths=(4, 4), + physical_size_fn=lambda _lengths: 12, + ), + ), + data_parallel_size=1, + ) + data = BatchedDataDict( + { + "input_ids": torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]]), + "input_lengths": torch.tensor([3, 2]), + "sample_mask": torch.tensor([1, 0]), + "batch_item_id": torch.tensor([10, 11]), + "student_semantic_regions": [ + ((10, "assistant", "eot", 1, 3),), + ((11, "assistant", "eot", 1, 2),), + ], + } + ) + data.lockstep_packing_plan = plan + data.lockstep_side_id = "student" + data.lockstep_bin_indices = (0,) + + with ( + patch( + "nemo_rl.models.megatron.data.get_context_parallel_world_size", + return_value=1, + ), + patch("nemo_rl.models.megatron.data.get_context_parallel_rank", return_value=0), + ): + processed = process_microbatch( + data, + seq_length_key="input_lengths", + pad_individual_seqs_to_multiple_of=4, + pad_packed_seq_to_multiple_of=4, + pack_sequences=True, + ) + + assert processed.input_ids.shape == (1, 12) + assert processed.packed_seq_params.cu_seqlens_q.tolist() == [0, 3, 5] + assert processed.packed_seq_params.cu_seqlens_q_padded.tolist() == [0, 4, 12] + assert processed.padding_mask.tolist() == [ + [False, False, False, True, True, True, True, True, True, True, True, True] + ] + assert data["student_semantic_regions"][0][0][0] == 10 + + data["input_lengths"] = torch.tensor([2, 2]) + with pytest.raises(ValueError, match="input lengths drifted"): + process_microbatch( + data, + seq_length_key="input_lengths", + pad_individual_seqs_to_multiple_of=4, + pad_packed_seq_to_multiple_of=4, + pack_sequences=True, + ) + + GET_PACK_SEQUENCE_PARAMETERS_TEST_ACTOR_FQN = f"{GetPackSequenceParametersTestActor.__module__}.GetPackSequenceParametersTestActor" diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 93c1db0388a..31e99ad6624 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -202,6 +202,36 @@ def test_model_name_is_hf_model(self, tmp_path): assert pretrained_path == f"{tmp_path}/meta-llama/Llama-3.2-1B" assert pt_checkpoint_exists is False + def test_megatron_fsdp_uses_local_hf_snapshot(self, monkeypatch): + """Megatron-FSDP initializes from HF instead of a torch_dist cache.""" + from nemo_rl.models.megatron.setup import validate_model_paths + + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + + config = { + "model_name": "meta-llama/Llama-3.1-8B", + "megatron_cfg": { + "distributed_data_parallel_config": { + "use_megatron_fsdp": True, + } + }, + } + + with patch( + "huggingface_hub.snapshot_download", return_value="/hf/snapshot" + ) as mock_snapshot_download: + hf_model_name, pretrained_path, pt_checkpoint_exists = validate_model_paths( + config + ) + + mock_snapshot_download.assert_called_once_with( + repo_id="meta-llama/Llama-3.1-8B", + local_files_only=False, + ) + assert hf_model_name == "meta-llama/Llama-3.1-8B" + assert pretrained_path == "/hf/snapshot" + assert pt_checkpoint_exists is True + def test_model_name_is_local_path(self, tmp_path): """Test with a local path as model name.""" from nemo_rl.models.megatron.setup import validate_model_paths @@ -766,8 +796,8 @@ def test_basic_parallelism_config(self): assert model_cfg.sequence_parallel is True assert model_cfg.context_parallel_size == 1 - def test_context_parallel_requires_sequence_packing(self): - """Test that context parallelism > 1 requires sequence packing.""" + def test_context_parallel_without_sequence_packing(self): + """Context parallelism > 1 is accepted with sequence packing disabled.""" from nemo_rl.models.megatron.setup import _apply_parallelism_config model_cfg = MagicMock() @@ -783,10 +813,9 @@ def test_context_parallel_requires_sequence_packing(self): "sequence_packing": {"enabled": False}, } - with pytest.raises(AssertionError) as exc_info: - _apply_parallelism_config(model_cfg, config) + _apply_parallelism_config(model_cfg, config) - assert "Sequence Packing must be enabled" in str(exc_info.value) + assert model_cfg.context_parallel_size == 2 def test_context_parallel_with_sequence_packing(self): """Test context parallelism with sequence packing enabled.""" @@ -2895,6 +2924,19 @@ def test_basic_checkpoint_config(self, tmp_path): assert checkpoint_config.fully_parallel_load is True assert checkpoint_config.load_rng is False + def test_megatron_fsdp_uses_fsdp_dtensor_format(self, tmp_path): + """Megatron-FSDP saves and resumes through its required DCP format.""" + from nemo_rl.models.megatron.setup import _create_checkpoint_config + + checkpoint_config = _create_checkpoint_config( + str(tmp_path / "pretrained"), + str(tmp_path / "weights"), + str(tmp_path / "optimizer"), + use_megatron_fsdp=True, + ) + + assert checkpoint_config.ckpt_format == "fsdp_dtensor" + def test_missing_ckpt_cfg_defaults_to_sync_save(self, tmp_path): """An absent checkpoint block keeps Megatron Bridge's default (sync save). @@ -3255,7 +3297,11 @@ def _megatron_config(**megatron_overrides: Any) -> dict[str, Any]: "train_iters": 10, } megatron_cfg.update(megatron_overrides) - return {"megatron_cfg": megatron_cfg, "train_global_batch_size": 8} + return { + "megatron_cfg": megatron_cfg, + "tokenizer": {"name": "pinned-tokenizer-snapshot"}, + "train_global_batch_size": 8, + } def _dist_config_passed_to_container(config: dict[str, Any]) -> Any: @@ -3292,6 +3338,33 @@ def _dist_config_passed_to_container(config: dict[str, Any]) -> Any: class TestCreateMegatronConfigGlooProcessGroups: """Tests for use_gloo_process_groups plumbing in _create_megatron_config.""" + def test_backend_tokenizer_uses_configured_tokenizer_snapshot(self): + """Backend tokenization must not silently fall back to the model id.""" + from nemo_rl.models.megatron.setup import _create_megatron_config + + config = _megatron_config() + with ( + patch("nemo_rl.models.megatron.setup.ConfigContainer"), + patch("nemo_rl.models.megatron.setup.TrainingConfig"), + patch("nemo_rl.models.megatron.setup.OptimizerConfig"), + patch("nemo_rl.models.megatron.setup.DistributedDataParallelConfig"), + patch("nemo_rl.models.megatron.setup.SchedulerConfig"), + patch("nemo_rl.models.megatron.setup.TokenizerConfig") as mock_tokenizer, + patch("nemo_rl.models.megatron.setup.LoggerConfig"), + ): + _create_megatron_config( + model_cfg=MagicMock(), + checkpoint_config=MagicMock(), + config=config, + hf_model_name="mutable-model-id", + dtype=torch.bfloat16, + ) + + mock_tokenizer.assert_called_once_with( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model="pinned-tokenizer-snapshot", + ) + @pytest.mark.parametrize("value", [True, False]) def test_explicit_value_is_forwarded(self, value): """An explicit use_gloo_process_groups is applied to the dist config.""" @@ -3312,6 +3385,111 @@ def test_absent_key_defers_to_bridge_default(self): == DistributedInitConfig().use_gloo_process_groups ) + def test_optimizer_dtype_strings_are_normalized(self): + """Precision-aware optimizer dtype overrides reach MCore as torch dtypes.""" + from nemo_rl.models.megatron.setup import _create_megatron_config + + config = _megatron_config() + config["megatron_cfg"]["optimizer"].update( + { + "params_dtype": "bfloat16", + "main_grads_dtype": "torch.bfloat16", + "main_params_dtype": "float32", + "exp_avg_dtype": "bfloat16", + "exp_avg_sq_dtype": "bfloat16", + } + ) + + with ( + patch("nemo_rl.models.megatron.setup.ConfigContainer"), + patch("nemo_rl.models.megatron.setup.TrainingConfig"), + patch( + "nemo_rl.models.megatron.setup.OptimizerConfig" + ) as mock_optimizer_config, + patch("nemo_rl.models.megatron.setup.DistributedDataParallelConfig"), + patch("nemo_rl.models.megatron.setup.SchedulerConfig"), + patch("nemo_rl.models.megatron.setup.TokenizerConfig"), + patch("nemo_rl.models.megatron.setup.LoggerConfig"), + ): + _create_megatron_config( + model_cfg=MagicMock(), + checkpoint_config=MagicMock(), + config=config, + hf_model_name="test-model", + dtype=torch.bfloat16, + ) + + optimizer_kwargs = mock_optimizer_config.call_args.kwargs + assert optimizer_kwargs["params_dtype"] is torch.bfloat16 + assert optimizer_kwargs["main_grads_dtype"] is torch.bfloat16 + assert optimizer_kwargs["main_params_dtype"] is torch.float32 + assert optimizer_kwargs["exp_avg_dtype"] is torch.bfloat16 + assert optimizer_kwargs["exp_avg_sq_dtype"] is torch.bfloat16 + + def test_fp32_accumulating_reduce_scatter_is_forwarded(self): + """The optional mixed-precision reduce-scatter reaches MCore DDP.""" + from nemo_rl.models.megatron.setup import _create_megatron_config + + config = _megatron_config() + config["megatron_cfg"]["distributed_data_parallel_config"][ + "reduce_scatter_with_fp32_accumulation" + ] = True + + with ( + patch("nemo_rl.models.megatron.setup.ConfigContainer"), + patch("nemo_rl.models.megatron.setup.TrainingConfig"), + patch("nemo_rl.models.megatron.setup.OptimizerConfig"), + patch( + "nemo_rl.models.megatron.setup.DistributedDataParallelConfig" + ) as mock_ddp_config, + patch("nemo_rl.models.megatron.setup.SchedulerConfig"), + patch("nemo_rl.models.megatron.setup.TokenizerConfig"), + patch("nemo_rl.models.megatron.setup.LoggerConfig"), + ): + _create_megatron_config( + model_cfg=MagicMock(), + checkpoint_config=MagicMock(), + config=config, + hf_model_name="test-model", + dtype=torch.bfloat16, + ) + + assert ( + mock_ddp_config.call_args.kwargs["reduce_scatter_with_fp32_accumulation"] + is True + ) + + def test_megatron_fsdp_is_forwarded_to_dist_and_ddp(self): + """The NeMo RL flag must select Bridge's Megatron-FSDP wrapper.""" + from nemo_rl.models.megatron.setup import _create_megatron_config + + config = _megatron_config() + config["megatron_cfg"]["distributed_data_parallel_config"][ + "use_megatron_fsdp" + ] = True + + with ( + patch("nemo_rl.models.megatron.setup.ConfigContainer") as mock_container, + patch("nemo_rl.models.megatron.setup.TrainingConfig"), + patch("nemo_rl.models.megatron.setup.OptimizerConfig"), + patch( + "nemo_rl.models.megatron.setup.DistributedDataParallelConfig" + ) as mock_ddp, + patch("nemo_rl.models.megatron.setup.SchedulerConfig"), + patch("nemo_rl.models.megatron.setup.TokenizerConfig"), + patch("nemo_rl.models.megatron.setup.LoggerConfig"), + ): + _create_megatron_config( + model_cfg=MagicMock(), + checkpoint_config=MagicMock(), + config=config, + hf_model_name="test-model", + dtype=torch.bfloat16, + ) + + assert mock_container.call_args.kwargs["dist"].use_megatron_fsdp is True + assert mock_ddp.call_args.kwargs["use_megatron_fsdp"] is True + @pytest.mark.mcore class TestDistributedTimeoutMinutes: @@ -3464,7 +3642,11 @@ def _config(fp8_cfg: dict[str, Any] | None) -> dict[str, Any]: } if fp8_cfg is not None: megatron_cfg["fp8_cfg"] = fp8_cfg - return {"megatron_cfg": megatron_cfg, "train_global_batch_size": 8} + return { + "megatron_cfg": megatron_cfg, + "tokenizer": {"name": "pinned-tokenizer-snapshot"}, + "train_global_batch_size": 8, + } @staticmethod def _optimizer_passed_to_container( @@ -3556,6 +3738,7 @@ def _subconfig_kwargs(): }, "train_iters": 10, }, + "tokenizer": {"name": "pinned-tokenizer-snapshot"}, "train_global_batch_size": 8, } with ( @@ -3618,6 +3801,7 @@ def test_fraction_and_transfer_overlap_are_forwarded(self, transfer_overlap): }, "train_iters": 10, }, + "tokenizer": {"name": "pinned-tokenizer-snapshot"}, "train_global_batch_size": 8, } @@ -3718,6 +3902,9 @@ def test_generation_refit_environment(self, backend, colocated, expected_cumem): ) as mock_setup_model_config: mock_megatron_cfg = MagicMock() mock_megatron_cfg.model.vocab_size = 32000 + mock_megatron_cfg.model.make_vocab_size_divisible_by = 128 + mock_megatron_cfg.model.tensor_model_parallel_size = 2 + mock_megatron_cfg.model.should_pad_vocab = True mock_setup_model_config.return_value = (mock_megatron_cfg, MagicMock()) with patch( @@ -3737,6 +3924,88 @@ def test_generation_refit_environment(self, backend, colocated, expected_cumem): assert runtime_config.offload_optimizer_for_refit is True assert os.environ.get("NCCL_CUMEM_ENABLE") == expected_cumem + def test_unpadded_model_uses_raw_vocab_size(self): + """Runtime output geometry must match an unpadded model head.""" + from nemo_rl.models.megatron.setup import validate_and_set_config + + config = { + "precision": "bfloat16", + "megatron_cfg": { + "optimizer": {"optimizer_cpu_offload": False}, + "tensor_model_parallel_size": 2, + }, + "offload_optimizer_for_logprob": False, + "offload_optimizer_for_refit": True, + } + mock_megatron_cfg = MagicMock() + mock_megatron_cfg.model.vocab_size = 151936 + mock_megatron_cfg.model.make_vocab_size_divisible_by = 128 + mock_megatron_cfg.model.tensor_model_parallel_size = 2 + mock_megatron_cfg.model.should_pad_vocab = False + + with ( + patch( + "nemo_rl.models.megatron.setup.setup_model_config", + return_value=(mock_megatron_cfg, mock_megatron_cfg.model), + ), + patch( + "nemo_rl.models.megatron.setup.calculate_padded_vocab_size", + return_value=152064, + ) as mock_calculate_padded_vocab_size, + ): + runtime_config = validate_and_set_config( + config=config, + rank=0, + hf_model_name="test-model", + pretrained_path="/path/to/model", + weights_path=None, + optimizer_path=None, + ) + + mock_calculate_padded_vocab_size.assert_not_called() + assert runtime_config.final_padded_vocab_size == 151936 + + def test_padded_model_uses_calculated_vocab_size(self): + """Runtime output geometry must include requested vocabulary padding.""" + from nemo_rl.models.megatron.setup import validate_and_set_config + + config = { + "precision": "bfloat16", + "megatron_cfg": { + "optimizer": {"optimizer_cpu_offload": False}, + "tensor_model_parallel_size": 2, + }, + "offload_optimizer_for_logprob": False, + "offload_optimizer_for_refit": True, + } + mock_megatron_cfg = MagicMock() + mock_megatron_cfg.model.vocab_size = 151936 + mock_megatron_cfg.model.make_vocab_size_divisible_by = 128 + mock_megatron_cfg.model.tensor_model_parallel_size = 2 + mock_megatron_cfg.model.should_pad_vocab = True + + with ( + patch( + "nemo_rl.models.megatron.setup.setup_model_config", + return_value=(mock_megatron_cfg, mock_megatron_cfg.model), + ), + patch( + "nemo_rl.models.megatron.setup.calculate_padded_vocab_size", + return_value=152064, + ) as mock_calculate_padded_vocab_size, + ): + runtime_config = validate_and_set_config( + config=config, + rank=0, + hf_model_name="test-model", + pretrained_path="/path/to/model", + weights_path=None, + optimizer_path=None, + ) + + mock_calculate_padded_vocab_size.assert_called_once_with(151936, 128, 2) + assert runtime_config.final_padded_vocab_size == 152064 + @pytest.mark.mcore class TestRuntimeConfigNamedTuple: @@ -4151,6 +4420,40 @@ def test_model_overrides_are_finalized_and_serialized(self, tmp_path, request): } assert model_cfg.masked_softmax_fusion is True + def test_megatron_fsdp_builds_model_config_from_hf(self, request): + """Direct HF initialization avoids reading a torch_dist run_config.""" + from nemo_rl.models.megatron.setup import setup_model_config + + self._apply_patches(request) + + mock_provider = MagicMock() + mock_provider.to_megatron_provider.return_value = self._make_model_cfg_mock() + config = { + "tokenizer": {"name": "pinned-tokenizer-snapshot"}, + "megatron_cfg": { + "distributed_data_parallel_config": { + "use_megatron_fsdp": True, + } + }, + } + + with ( + patch("transformers.AutoConfig.from_pretrained") as mock_ac, + patch("nemo_rl.models.megatron.setup.AutoBridge") as mock_ab, + ): + mock_ab.from_hf_config.return_value = mock_provider + setup_model_config( + config, + rank=0, + dtype=torch.bfloat16, + hf_model_name="test-model", + pretrained_path="/hf/snapshot", + ) + + mock_ac.assert_called_once_with("test-model", trust_remote_code=True) + mock_ab.from_hf_config.assert_called_once_with(mock_ac.return_value) + mock_provider.to_megatron_provider.return_value.finalize.assert_called_once_with() + def test_megatron_lm_no_overrides_calls_autoconfig_without_extra_kwargs( self, request ): @@ -4413,6 +4716,57 @@ def test_force_reconvert_from_hf_when_checkpoint_exists( mock_ps.destroy_model_parallel.assert_not_called() +@pytest.mark.mcore +class TestFreezeParameterPatternsPreWrapHook: + """Tests for the MCore equivalent of policy.freeze_parameter_patterns.""" + + def test_freezes_substring_matches_across_model_chunks(self, capsys): + from nemo_rl.models.megatron.setup import ( + _create_freeze_parameter_patterns_pre_wrap_hook, + ) + + first_chunk = torch.nn.Module() + first_chunk.mtp = torch.nn.Linear(2, 2) + first_chunk.layer0 = torch.nn.Linear(2, 2) + # An already-frozen match is counted but must not be counted as newly frozen. + first_chunk.mtp.bias.requires_grad_(False) + + second_chunk = torch.nn.Module() + second_chunk.mixer = torch.nn.Module() + second_chunk.mixer.gate = torch.nn.Linear(2, 2) + second_chunk.keep = torch.nn.Linear(2, 2) + + model = [first_chunk, second_chunk] + hook = _create_freeze_parameter_patterns_pre_wrap_hook( + ["mtp", "mixer.gate", "layer[0]"] + ) + + result = hook(model) + + assert result is model + assert all(not param.requires_grad for param in first_chunk.mtp.parameters()) + assert all( + not param.requires_grad for param in second_chunk.mixer.gate.parameters() + ) + # Patterns are literal substrings, not regular expressions. + assert all(param.requires_grad for param in first_chunk.layer0.parameters()) + assert all(param.requires_grad for param in second_chunk.keep.parameters()) + assert "matched 4 parameters; newly froze 3" in capsys.readouterr().out + + def test_returns_single_model_unchanged_when_nothing_matches(self): + from nemo_rl.models.megatron.setup import ( + _create_freeze_parameter_patterns_pre_wrap_hook, + ) + + model = torch.nn.Linear(2, 2) + hook = _create_freeze_parameter_patterns_pre_wrap_hook(["does.not.match"]) + + result = hook(model) + + assert result is model + assert all(param.requires_grad for param in model.parameters()) + + @pytest.mark.mcore class TestSetupModelAndOptimizer: """Tests for setup_model_and_optimizer function.""" @@ -4481,22 +4835,40 @@ def test_setup_with_param_sync_and_frozen_moe_router( mock_checkpoint_exists.return_value = False policy_cfg = { + "freeze_parameter_patterns": ["mtp"], "megatron_cfg": { "freeze_moe_router": True, # Enable MoE router freezing - } + }, } + def additional_hook(model): + return model + result = setup_model_and_optimizer( policy_cfg=policy_cfg, megatron_cfg=mock_megatron_cfg, load_optimizer=True, + additional_pre_wrap_hooks=[additional_hook], ) # Verify get_model was called (the mixed_precision_wrapper should be CustomFloat16Module) mock_get_model.assert_called_once() call_kwargs = mock_get_model.call_args[1] - # Check that pre_wrap_hook is not empty when freeze_moe_router is True - assert len(call_kwargs.get("pre_wrap_hook", [])) > 0 + assert ( + call_kwargs["use_megatron_fsdp"] is mock_megatron_cfg.dist.use_megatron_fsdp + ) + hooks = call_kwargs.get("pre_wrap_hook", []) + assert [hook.__name__ for hook in hooks] == [ + "freeze_moe_router", + "additional_hook", + "freeze_parameter_patterns", + ] + # Bridge composes these hooks by passing one hook's return value into the + # next, so each hook must preserve and return the model object. + hooked_model = mock_model + for hook in hooks: + hooked_model = hook(hooked_model) + assert hooked_model is mock_model assert result.param_sync_func == mock_model_chunk.start_param_sync @@ -4649,6 +5021,7 @@ def test_basic_finalize_setup( mock_auto_bridge.from_hf_pretrained.return_value = mock_bridge config = { + "tokenizer": {"name": "pinned-tokenizer-snapshot"}, "megatron_cfg": { "tensor_model_parallel_size": 2, "optimizer": { @@ -4657,7 +5030,7 @@ def test_basic_finalize_setup( "distributed_data_parallel_config": { "overlap_param_gather": False, }, - } + }, } result = finalize_megatron_setup( @@ -4681,6 +5054,8 @@ def test_basic_finalize_setup( mock_get_model_config.assert_called_once_with(mock_model) assert mock_update_model_config.call_args.args[1] is runtime_model_config mock_build_tokenizer.assert_called_once() + tokenizer_config = mock_build_tokenizer.call_args.args[0] + assert tokenizer_config.tokenizer_model == "pinned-tokenizer-snapshot" mock_auto_bridge.from_hf_pretrained.assert_called_once_with( "test-model", trust_remote_code=True ) diff --git a/tests/unit/models/megatron/test_train.py b/tests/unit/models/megatron/test_train.py index 62ae69acb79..8e08da9e2a2 100644 --- a/tests/unit/models/megatron/test_train.py +++ b/tests/unit/models/megatron/test_train.py @@ -32,6 +32,7 @@ from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams from nemo_rl.algorithms.loss.interfaces import LossInputType +from nemo_rl.distributed.batched_data_dict import BatchedDataDict pytestmark = pytest.mark.mcore @@ -327,6 +328,30 @@ def __init__(self): mock_scatter.assert_not_called() assert result is padding_mask + def test_model_forward_passes_cp_local_router_padding_mask(self): + """Packed physical/masked tokens are excluded by GPT's MoE router.""" + from nemo_rl.models.megatron.train import model_forward + + mock_model = MagicMock(return_value=torch.randn(1, 6, 8)) + mock_model.config = SimpleNamespace(sequence_parallel=False) + mock_data_dict = MagicMock() + mock_data_dict.get_multimodal_dict.return_value = {} + padding_mask = torch.tensor( + [[False, True, True, True, True, True]], dtype=torch.bool + ) + + model_forward( + model=mock_model, + data_dict=mock_data_dict, + input_ids_cp_sharded=torch.tensor([[1, 0, 4, 0, 0, 0]]), + position_ids=None, + attention_mask=None, + packed_seq_params=MagicMock(), + padding_mask=padding_mask, + ) + + assert mock_model.call_args.kwargs["padding_mask"] is padding_mask + def test_model_forward_with_defer_fp32_logits(self): """Test model_forward passes fp32_output when defer_fp32_logits is True.""" from nemo_rl.models.megatron.train import model_forward @@ -1160,6 +1185,7 @@ def test_megatron_forward_backward_calls_forward_backward_func(self, mock_get_fb cfg = {"sequence_packing": {"enabled": False}} post_processor = LossPostProcessor(loss_fn=mock_loss_fn, cfg=cfg) + global_valid_chunks_by_idx = {0: torch.tensor(37)} megatron_forward_backward( model=mock_model, data_iterator=iter([]), @@ -1167,6 +1193,7 @@ def test_megatron_forward_backward_calls_forward_backward_func(self, mock_get_fb seq_length=128, mbs=2, post_processing_fn=post_processor, + global_valid_chunks_by_idx=global_valid_chunks_by_idx, ) mock_get_fb.assert_called_once() @@ -1177,6 +1204,10 @@ def test_megatron_forward_backward_calls_forward_backward_func(self, mock_get_fb assert call_kwargs["num_microbatches"] == 4 assert call_kwargs["seq_length"] == 128 assert call_kwargs["micro_batch_size"] == 2 + assert ( + call_kwargs["forward_step_func"].keywords["global_valid_chunks_by_idx"] + is global_valid_chunks_by_idx + ) @patch("nemo_rl.models.megatron.train.get_forward_backward_func") def test_megatron_forward_backward_forward_only(self, mock_get_fb): @@ -1526,9 +1557,13 @@ def test_loss_post_processor_no_packing( self, mock_cp_size, mock_cp_grp, mock_tp_grp, mock_tp_rank ): """Test LossPostProcessor without sequence packing.""" + from nemo_rl.algorithms.loss.csd import CSDLossFn from nemo_rl.models.megatron.train import LossPostProcessor - mock_loss_fn = MagicMock(return_value=(torch.tensor(0.5), {"loss": 0.5})) + # Only CSD losses receive the global alignment-chunk counts. + mock_loss_fn = MagicMock( + spec=CSDLossFn, return_value=(torch.tensor(0.5), {"loss": 0.5}) + ) mock_loss_fn.input_type = LossInputType.LOGIT cfg = {"sequence_packing": {"enabled": False}} @@ -1538,11 +1573,13 @@ def test_loss_post_processor_no_packing( mock_tp_grp.return_value = MagicMock() mock_cp_grp.return_value = MagicMock() + global_valid_chunks_by_idx = {0: torch.tensor(37)} wrapped_fn = processor( data_dict=MagicMock(), packed_seq_params=None, global_valid_seqs=torch.tensor(10), global_valid_toks=torch.tensor(100), + global_valid_chunks_by_idx=global_valid_chunks_by_idx, ) # Call the wrapped function @@ -1552,6 +1589,10 @@ def test_loss_post_processor_no_packing( assert torch.isclose(loss, torch.tensor(0.5)) assert isinstance(metrics, dict) assert len(metrics) == 1 and metrics["loss"] == 0.5 + assert ( + mock_loss_fn.call_args.kwargs["global_valid_chunks_by_idx"] + is global_valid_chunks_by_idx + ) @patch( "nemo_rl.models.megatron.train.get_tensor_model_parallel_rank", return_value=0 @@ -1587,6 +1628,51 @@ def test_loss_post_processor_with_cp_normalize( # Loss should be scaled by num_microbatches / (cp_size * cp_size) = 4 / (2 * 2) = 1.0 assert torch.isclose(loss, torch.tensor(1.0)) + @patch( + "nemo_rl.models.megatron.train.get_tensor_model_parallel_rank", return_value=0 + ) + @patch("nemo_rl.models.megatron.train.get_tensor_model_parallel_group") + @patch("nemo_rl.models.megatron.train.get_context_parallel_group") + @patch( + "nemo_rl.models.megatron.train.get_context_parallel_world_size", return_value=2 + ) + def test_cp_partial_loss_keeps_only_the_schedule_compensation( + self, mock_cp_size, mock_cp_grp, mock_tp_grp, mock_tp_rank + ): + """A CP-partial loss is not divided by CP; Megatron's own scaling is undone.""" + from nemo_rl.models.megatron.train import LossPostProcessor + + mock_loss_fn = MagicMock(return_value=(torch.tensor(1.0), {})) + mock_loss_fn.input_type = LossInputType.LOGIT + processor = LossPostProcessor( + loss_fn=mock_loss_fn, + cfg={"sequence_packing": {"enabled": False}}, + num_microbatches=4, + cp_normalize=False, + ) + mock_tp_grp.return_value = MagicMock() + mock_cp_grp.return_value = MagicMock() + + loss, _ = processor(data_dict=MagicMock())(torch.randn(2, 10, 100)) + + # Only num_microbatches / cp_size = 4 / 2, which cancels the schedule's + # cp_size / num_microbatches; summed over CP the partials give the loss. + assert torch.isclose(loss, torch.tensor(2.0)) + + def test_only_owner_local_csd_returns_a_cp_partial_loss(self): + from nemo_rl.algorithms.loss.csd import CSDLossConfig, CSDLossFn + from nemo_rl.models.megatron.train import ( + loss_returns_context_parallel_partial, + ) + + assert loss_returns_context_parallel_partial( + CSDLossFn(CSDLossConfig(student_vocab_size=8, owner_local=True)) + ) + assert not loss_returns_context_parallel_partial( + CSDLossFn(CSDLossConfig(student_vocab_size=8)) + ) + assert not loss_returns_context_parallel_partial(MagicMock()) + @patch( "nemo_rl.models.megatron.train.get_tensor_model_parallel_rank", return_value=0 ) @@ -1620,6 +1706,52 @@ def test_loss_post_processor_with_packing( # Verify SequencePackingLossWrapper was called mock_wrapper.assert_called_once() + def test_csd_packing_selects_axis_aware_wrapper_even_when_fusion_enabled( + self, + ): + import nemo_rl.models.megatron.train as megatron_train + from nemo_rl.algorithms.loss.csd import CSDLossFn + + packed_params = MagicMock( + cu_seqlens_q=torch.tensor([0, 3], dtype=torch.int32), + cu_seqlens_q_padded=torch.tensor([0, 4], dtype=torch.int32), + ) + with ( + patch.object( + megatron_train, "get_tensor_model_parallel_rank", return_value=0 + ), + patch.object( + megatron_train, + "get_tensor_model_parallel_group", + return_value=MagicMock(), + ), + patch.object( + megatron_train, + "get_context_parallel_group", + return_value=MagicMock(), + ), + patch.object( + megatron_train, "CSDSequencePackingLossWrapper" + ) as mock_csd_wrapper, + patch.object( + megatron_train, "SequencePackingFusionLossWrapper" + ) as mock_fused_wrapper, + patch.object( + megatron_train, "SequencePackingLossWrapper" + ) as mock_generic_wrapper, + ): + processor = megatron_train.LossPostProcessor( + loss_fn=object.__new__(CSDLossFn), + cfg={"sequence_packing": {"enabled": True, "fuse_loss": True}}, + cp_normalize=False, + ) + + processor(data_dict=MagicMock(), packed_seq_params=packed_params) + + mock_csd_wrapper.assert_called_once() + mock_fused_wrapper.assert_not_called() + mock_generic_wrapper.assert_not_called() + def test_direct_model_loss_normalizes_target_aligned_tokens_and_schedule_scaling(): loss, _ = _run_direct_model_loss() @@ -2211,6 +2343,548 @@ def fake_allgather(local_tensor, group, seq_dim): torch.testing.assert_close(result[key], expected) +class TestSparseTeacherLogitsPostProcessor: + """Tests for the CP-local sparse teacher row export.""" + + @staticmethod + def _sparse_outputs(local_seq_len: int, k: int, force_width: int = 2): + return ( + torch.randn(1, local_seq_len, k), + torch.randint(0, 32, (1, local_seq_len, k)), + torch.randn(1, local_seq_len), + torch.randint(0, 32, (1, local_seq_len, 1)), + torch.randn(1, local_seq_len, force_width), + torch.randint(0, 32, (1, local_seq_len, force_width)), + torch.zeros(1, local_seq_len, force_width, dtype=torch.bool), + ) + + @patch("nemo_rl.models.megatron.train.allgather_cp_sharded_tensor") + @patch("nemo_rl.models.megatron.train.distributed_vocab_topk_logz_force") + @patch( + "nemo_rl.models.megatron.train.get_tensor_model_parallel_rank", + return_value=1, + ) + @patch("nemo_rl.models.megatron.train.get_tensor_model_parallel_group") + @patch("torch.distributed.get_rank") + @patch("nemo_rl.models.megatron.train.get_context_parallel_group") + @patch("nemo_rl.models.megatron.train.get_context_parallel_world_size") + def test_force_ids_follow_load_balanced_cp_rows_without_gather( + self, + mock_cp_world, + mock_cp_group, + mock_dist_rank, + mock_tp_group, + mock_tp_rank, + mock_sparse_topk, + mock_cp_allgather, + ): + """Forced-label rows use the logits' load-balanced CP row order.""" + from nemo_rl.models.megatron.train import SparseTeacherLogitsPostProcessor + + mock_cp_group.return_value = MagicMock() + tp_group = MagicMock() + mock_tp_group.return_value = tp_group + local_seq_len = 4 + k = 3 + for cp_size, cp_rank, full_seq_len, expected_positions in ( + (2, 1, 8, [2, 3, 4, 5]), + (4, 2, 16, [4, 5, 10, 11]), + ): + mock_cp_world.return_value = cp_size + mock_dist_rank.return_value = cp_rank + full_force_ids = torch.arange(full_seq_len * 2).reshape(1, full_seq_len, 2) + mock_sparse_topk.return_value = self._sparse_outputs(local_seq_len, k) + processor = SparseTeacherLogitsPostProcessor( + cfg={ + "sequence_packing": {"enabled": False}, + "logprob_chunk_size": 2, + }, + k=k, + temperature=0.5, + real_vocab_size=30, + noise_filter_k=5, + ) + + wrapped = processor( + data_dict={"force_include_token_ids": full_force_ids}, + cu_seqlens_padded=None, + ) + loss, result = wrapped(torch.randn(1, local_seq_len, 16)) + + call = mock_sparse_topk.call_args + assert torch.equal(call.args[1], full_force_ids[:, expected_positions]) + assert call.args[2] == k + assert call.args[3] is tp_group + assert call.kwargs == { + "vocab_start_index": 16, + "vocab_end_index": 32, + "real_vocab_size": 30, + "temperature": 0.5, + "noise_filter_k": 5, + "chunk_size": 2, + } + mock_cp_allgather.assert_not_called() + assert loss.item() == 0.0 + assert set(result) == { + "sparse_topk_logits", + "sparse_topk_indices", + "sparse_log_z", + "sparse_natural_tail_indices", + "sparse_forced_logits", + "sparse_forced_indices", + "sparse_forced_in_topk", + } + assert result["sparse_topk_logits"].shape == (1, local_seq_len, k) + + @patch("nemo_rl.models.megatron.train.distributed_vocab_topk_logz_force") + @patch( + "nemo_rl.models.megatron.train.get_tensor_model_parallel_rank", + return_value=0, + ) + @patch("nemo_rl.models.megatron.train.get_tensor_model_parallel_group") + @patch( + "nemo_rl.models.megatron.train.get_context_parallel_world_size", + return_value=1, + ) + def test_unset_logprob_chunk_size_uses_bounded_row_tiles( + self, _mock_cp_world, _mock_tp_group, _mock_tp_rank, mock_sparse_topk + ): + from nemo_rl.models.megatron import train as megatron_train + + mock_sparse_topk.return_value = self._sparse_outputs(4, 2) + processor = megatron_train.SparseTeacherLogitsPostProcessor( + cfg={"sequence_packing": {"enabled": False}}, + k=2, + temperature=1.0, + real_vocab_size=8, + noise_filter_k=0, + ) + wrapped = processor( + data_dict={"force_include_token_ids": torch.full((1, 4, 2), -1)}, + cu_seqlens_padded=None, + ) + wrapped(torch.randn(1, 4, 8)) + + assert ( + mock_sparse_topk.call_args.kwargs["chunk_size"] + == megatron_train.SPARSE_TEACHER_ROW_CHUNK_SIZE + ) + + @patch("nemo_rl.models.megatron.train.get_context_parallel_group") + @patch( + "nemo_rl.models.megatron.train.get_context_parallel_world_size", + return_value=2, + ) + @patch("torch.distributed.get_rank", return_value=0) + def test_force_ids_not_divisible_by_twice_cp_raise( + self, _mock_rank, _mock_cp_world, _mock_cp_group + ): + from nemo_rl.models.megatron.train import SparseTeacherLogitsPostProcessor + + processor = SparseTeacherLogitsPostProcessor( + cfg={"sequence_packing": {"enabled": False}}, + k=2, + temperature=1.0, + real_vocab_size=8, + noise_filter_k=0, + ) + with pytest.raises(ValueError, match="multiple of 2\\*cp_size=4"): + processor( + data_dict={"force_include_token_ids": torch.full((1, 6, 2), -1)}, + cu_seqlens_padded=None, + ) + + @patch( + "nemo_rl.models.megatron.train.get_context_parallel_world_size", + return_value=1, + ) + def test_force_rows_must_match_logit_rows(self, _mock_cp_world): + from nemo_rl.models.megatron.train import SparseTeacherLogitsPostProcessor + + processor = SparseTeacherLogitsPostProcessor( + cfg={"sequence_packing": {"enabled": False}}, + k=2, + temperature=1.0, + real_vocab_size=8, + noise_filter_k=0, + ) + wrapped = processor( + data_dict={"force_include_token_ids": torch.full((1, 6, 2), -1)}, + cu_seqlens_padded=None, + ) + with pytest.raises(ValueError, match="force=6, logits=8"): + wrapped(torch.randn(1, 8, 8)) + + def test_packing_and_missing_force_ids_raise(self): + from nemo_rl.models.megatron.train import SparseTeacherLogitsPostProcessor + + packed = SparseTeacherLogitsPostProcessor( + cfg={"sequence_packing": {"enabled": True}}, + k=2, + temperature=1.0, + real_vocab_size=8, + noise_filter_k=0, + ) + with pytest.raises(NotImplementedError, match="sequence packing"): + packed(data_dict={}, cu_seqlens_padded=None) + unpacked = SparseTeacherLogitsPostProcessor( + cfg={"sequence_packing": {"enabled": False}}, + k=2, + temperature=1.0, + real_vocab_size=8, + noise_filter_k=0, + ) + with pytest.raises(KeyError, match="force_include_token_ids"): + unpacked(data_dict={}, cu_seqlens_padded=None) + + @patch("nemo_rl.models.megatron.train.model_forward") + @patch("nemo_rl.models.megatron.train.apply_temperature_scaling") + def test_forward_dispatch_keeps_logits_unscaled( + self, mock_temp_scaling, mock_model_forward + ): + """The exporter applies the CSD temperature itself, never the sampler's.""" + from nemo_rl.models.megatron.data import ProcessedMicrobatch + from nemo_rl.models.megatron.train import ( + SparseTeacherLogitsPostProcessor, + forward_with_post_processing_fn, + ) + + mock_model_forward.return_value = torch.randn(1, 4, 8) + data_dict = MagicMock() + processed_mb = ProcessedMicrobatch( + data_dict=data_dict, + input_ids=torch.tensor([[1, 2, 3, 4]]), + input_ids_cp_sharded=torch.tensor([[1, 2, 3, 4]]), + attention_mask=None, + position_ids=None, + packed_seq_params=None, + cu_seqlens_padded=None, + original_seq_length=4, + ) + post_processor = SparseTeacherLogitsPostProcessor( + cfg={"sequence_packing": {"enabled": False}}, + k=2, + temperature=1.0, + real_vocab_size=8, + noise_filter_k=0, + ) + with patch.object( + SparseTeacherLogitsPostProcessor, "__call__", return_value=MagicMock() + ) as mock_call: + forward_with_post_processing_fn( + data_iterator=iter([processed_mb]), + model=MagicMock(), + post_processing_fn=post_processor, + sampling_params=TrainingSamplingParams(temperature=2.0), + ) + + mock_temp_scaling.assert_not_called() + mock_call.assert_called_once_with(data_dict=data_dict, cu_seqlens_padded=None) + + +class TestFullLogitsPostProcessor: + """Tests for CP-aware full-vocabulary teacher-logit export.""" + + def test_context_parallel_logits_are_relaid_to_contiguous_window(self): + import nemo_rl.models.megatron.train as megatron_train + + cp_group = MagicMock() + contiguous_logits = torch.randn(1, 2, 8, dtype=torch.float32) + with ( + patch.object( + megatron_train, "get_context_parallel_group", return_value=cp_group + ), + patch.object( + megatron_train, "get_context_parallel_world_size", return_value=2 + ), + patch.object( + megatron_train, + "cp_load_balanced_to_contiguous", + return_value=contiguous_logits, + ) as mock_relayout, + ): + processor = megatron_train.FullLogitsPostProcessor( + cfg={ + "sequence_packing": {"enabled": False}, + "megatron_cfg": {"context_parallel_size": 2}, + } + ) + wrapped_fn = processor(data_dict=MagicMock(), cu_seqlens_padded=None) + + output_tensor = torch.randn(1, 2, 8, dtype=torch.bfloat16) + loss, result = wrapped_fn(output_tensor) + + assert loss.item() == 0.0 + assert result["full_logits"] is contiguous_logits + relayout_args = mock_relayout.call_args + assert torch.equal(relayout_args.args[0], output_tensor.float()) + assert relayout_args.kwargs == {"cp_group": cp_group, "seq_dim": 1} + + def test_packed_cp1_restores_dense_logical_rows(self): + from megatron.core.packed_seq_params import PackedSeqParams + + import nemo_rl.models.megatron.train as megatron_train + + data = BatchedDataDict( + { + "input_ids": torch.zeros(2, 6, dtype=torch.long), + "input_lengths": torch.tensor([3, 2]), + } + ) + raw_cu = torch.tensor([0, 3, 5], dtype=torch.int32) + padded_cu = torch.tensor([0, 4, 8], dtype=torch.int32) + packed_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=raw_cu, + cu_seqlens_kv=raw_cu, + cu_seqlens_q_padded=padded_cu, + cu_seqlens_kv_padded=padded_cu, + ) + output = torch.tensor( + [ + [ + [1.0], + [2.0], + [3.0], + [99.0], + [10.0], + [11.0], + [98.0], + [97.0], + ] + ], + dtype=torch.bfloat16, + ) + + with patch.object( + megatron_train, "get_context_parallel_world_size", return_value=1 + ): + wrapped = megatron_train.FullLogitsPostProcessor( + cfg={ + "sequence_packing": {"enabled": True}, + "megatron_cfg": {"context_parallel_size": 1}, + } + )( + data_dict=data, + cu_seqlens_padded=padded_cu, + packed_seq_params=packed_params, + ) + loss, result = wrapped(output) + + assert loss.item() == 0.0 + restored = result["full_logits"] + assert restored.shape == (2, 6, 1) + assert restored[0, :, 0].tolist() == [1.0, 2.0, 3.0, 0.0, 0.0, 0.0] + assert restored[1, :, 0].tolist() == [10.0, 11.0, 0.0, 0.0, 0.0, 0.0] + + def test_packed_cp1_streams_valid_prefixes_without_retaining_tensors(self): + from megatron.core.packed_seq_params import PackedSeqParams + + import nemo_rl.models.megatron.train as megatron_train + + data = BatchedDataDict( + { + "input_ids": torch.zeros(2, 6, dtype=torch.long), + "input_lengths": torch.tensor([3, 2]), + } + ) + raw_cu = torch.tensor([0, 3, 5], dtype=torch.int32) + padded_cu = torch.tensor([0, 4, 8], dtype=torch.int32) + packed_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=raw_cu, + cu_seqlens_kv=raw_cu, + cu_seqlens_q_padded=padded_cu, + cu_seqlens_kv_padded=padded_cu, + ) + output = torch.tensor( + [ + [ + [1.0], + [2.0], + [3.0], + [99.0], + [10.0], + [11.0], + [98.0], + [97.0], + ] + ], + dtype=torch.bfloat16, + ) + storage = torch.full((5, 1), -1.0, dtype=torch.float32) + new_zeros_shapes = [] + original_new_zeros = torch.Tensor.new_zeros + + def record_new_zeros(tensor, size, *args, **kwargs): + new_zeros_shapes.append(tuple(size)) + return original_new_zeros(tensor, size, *args, **kwargs) + + with ( + patch.object( + megatron_train, "get_context_parallel_world_size", return_value=1 + ), + patch.object(torch.Tensor, "new_zeros", record_new_zeros), + ): + processor = megatron_train.FullLogitsPostProcessor( + cfg={ + "sequence_packing": {"enabled": True}, + "megatron_cfg": {"context_parallel_size": 1}, + }, + packed_output_storage=storage, + ) + wrapped = processor( + data_dict=data, + cu_seqlens_padded=padded_cu, + packed_seq_params=packed_params, + ) + _, result = wrapped(output) + + metadata = result["full_logits"] + assert isinstance(metadata, megatron_train.StreamedFullLogitsMetadata) + assert metadata.batch_size == 2 + assert metadata.logical_seq_len == 6 + assert metadata.storage_token_offset == 0 + assert metadata.stored_seq_lengths == (3, 2) + assert processor.packed_output_token_cursor == 5 + assert storage[:, 0].tolist() == [1.0, 2.0, 3.0, 10.0, 11.0] + assert new_zeros_shapes == [()] + assert all( + not isinstance(getattr(metadata, field_name), torch.Tensor) + for field_name in metadata.__dataclass_fields__ + ) + + def test_packed_cp2_uses_per_sequence_head_tail_boundaries(self): + from megatron.core.packed_seq_params import PackedSeqParams + + import nemo_rl.models.megatron.train as megatron_train + + cp_group = MagicMock() + data = BatchedDataDict( + { + "input_ids": torch.zeros(2, 8, dtype=torch.long), + "input_lengths": torch.tensor([3, 5]), + } + ) + raw_cu = torch.tensor([0, 3, 8], dtype=torch.int32) + padded_cu = torch.tensor([0, 4, 12], dtype=torch.int32) + packed_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=raw_cu, + cu_seqlens_kv=raw_cu, + cu_seqlens_q_padded=padded_cu, + cu_seqlens_kv_padded=padded_cu, + ) + local_output = torch.arange(6, dtype=torch.float32).reshape(1, 6, 1) + gathered = [ + torch.tensor([[[1.0], [2.0], [3.0], [99.0]]]), + torch.tensor( + [[[10.0], [11.0], [12.0], [13.0], [14.0], [98.0], [97.0], [96.0]]] + ), + ] + + with ( + patch.object( + megatron_train, "get_context_parallel_world_size", return_value=2 + ), + patch.object( + megatron_train, "get_context_parallel_group", return_value=cp_group + ), + patch.object(torch.distributed, "get_rank", return_value=0), + patch.object( + megatron_train, + "allgather_cp_sharded_tensor", + side_effect=gathered, + ) as mock_allgather, + ): + wrapped = megatron_train.FullLogitsPostProcessor( + cfg={ + "sequence_packing": {"enabled": True}, + "megatron_cfg": {"context_parallel_size": 2}, + } + )( + data_dict=data, + cu_seqlens_padded=padded_cu, + packed_seq_params=packed_params, + ) + _, result = wrapped(local_output) + + restored = result["full_logits"] + assert restored.shape == (2, 4, 1) + assert restored[0, :, 0].tolist() == [1.0, 2.0, 3.0, 0.0] + assert restored[1, :, 0].tolist() == [10.0, 11.0, 12.0, 13.0] + assert mock_allgather.call_count == 2 + assert mock_allgather.call_args_list[0].args[0].shape == (1, 2, 1) + assert mock_allgather.call_args_list[1].args[0].shape == (1, 4, 1) + + def test_packed_cp2_streams_zero_overlap_rows_and_keeps_collective_order(self): + from megatron.core.packed_seq_params import PackedSeqParams + + import nemo_rl.models.megatron.train as megatron_train + + cp_group = MagicMock() + data = BatchedDataDict( + { + "input_ids": torch.zeros(2, 8, dtype=torch.long), + "input_lengths": torch.tensor([3, 5]), + } + ) + raw_cu = torch.tensor([0, 3, 8], dtype=torch.int32) + padded_cu = torch.tensor([0, 4, 12], dtype=torch.int32) + packed_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=raw_cu, + cu_seqlens_kv=raw_cu, + cu_seqlens_q_padded=padded_cu, + cu_seqlens_kv_padded=padded_cu, + ) + local_output = torch.arange(6, dtype=torch.float32).reshape(1, 6, 1) + gathered = [ + torch.tensor([[[1.0], [2.0], [3.0], [99.0]]]), + torch.tensor( + [[[10.0], [11.0], [12.0], [13.0], [14.0], [98.0], [97.0], [96.0]]] + ), + ] + storage = torch.full((1, 1), -1.0, dtype=torch.float32) + + with ( + patch.object( + megatron_train, "get_context_parallel_world_size", return_value=2 + ), + patch.object( + megatron_train, "get_context_parallel_group", return_value=cp_group + ), + patch.object(torch.distributed, "get_rank", return_value=1), + patch.object( + megatron_train, + "allgather_cp_sharded_tensor", + side_effect=gathered, + ) as mock_allgather, + ): + processor = megatron_train.FullLogitsPostProcessor( + cfg={ + "sequence_packing": {"enabled": True}, + "megatron_cfg": {"context_parallel_size": 2}, + }, + packed_output_storage=storage, + ) + wrapped = processor( + data_dict=data, + cu_seqlens_padded=padded_cu, + packed_seq_params=packed_params, + ) + _, result = wrapped(local_output) + + metadata = result["full_logits"] + assert isinstance(metadata, megatron_train.StreamedFullLogitsMetadata) + assert metadata.stored_seq_lengths == (0, 1) + assert processor.packed_output_token_cursor == 1 + assert storage[:, 0].tolist() == [14.0] + # Every CP rank participates once per logical sample, even when this + # rank stores no tokens for the first row. + assert mock_allgather.call_count == 2 + + class TestAggregateTrainingStatistics: """Tests for aggregate_training_statistics function.""" diff --git a/tests/unit/models/policy/test_automodel_worker.py b/tests/unit/models/policy/test_automodel_worker.py index 941f6a82fd7..e36618786e1 100644 --- a/tests/unit/models/policy/test_automodel_worker.py +++ b/tests/unit/models/policy/test_automodel_worker.py @@ -1784,7 +1784,9 @@ def test_packed_dummy_batches_do_not_affect_metrics( ] assert counts == [2, 1] monkeypatch.setattr( - policy, "_shard_for_train", lambda data, batch_size: sharded_data + policy, + "_shard_for_train", + lambda data, batch_size, **_: sharded_data, ) policy.prepare_for_training() diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index d6515b9063d..bba6f8e2b3b 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -17,6 +17,7 @@ import os import tempfile import time +import weakref from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace @@ -2811,6 +2812,592 @@ def test_offload_after_refit_uses_standard_cleanup(monkeypatch): empty_cache.assert_not_called() +def test_full_logits_ipc_uses_valid_prefix_slab_and_preserves_logical_indices( + monkeypatch, +): + from megatron.core.packed_seq_params import PackedSeqParams + + from nemo_rl.models.policy.workers import megatron_policy_worker as mpw + + batch_sizes = [7, 2, 1] + batch_item_ids = [list(range(10, 17)), [20, 21], [30]] + input_lengths = [[1, 4, 2, 0, 3, 1, 2], [3, 1], [4]] + processed_microbatches = [] + for ids, lengths in zip(batch_item_ids, input_lengths, strict=True): + cumulative_lengths = [0] + for length in lengths: + cumulative_lengths.append(cumulative_lengths[-1] + length) + cu_seqlens = torch.tensor(cumulative_lengths, dtype=torch.int32) + packed_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens, + cu_seqlens_kv_padded=cu_seqlens, + ) + processed_microbatches.append( + SimpleNamespace( + data_dict=BatchedDataDict( + { + "batch_item_id": torch.tensor(ids), + "input_ids": torch.zeros(len(ids), 4, dtype=torch.long), + "input_lengths": torch.tensor(lengths), + } + ), + cu_seqlens_padded=cu_seqlens, + packed_seq_params=packed_params, + ) + ) + + def fake_get_microbatch_iterator(*_args, **_kwargs): + return iter(processed_microbatches), len(batch_sizes), 1, 4, 4 + + retained_outputs = [] + + def fake_megatron_forward_backward(**kwargs): + outputs = [] + for bin_index, processed_mb in enumerate(kwargs["data_iterator"]): + physical_tokens = int(processed_mb.cu_seqlens_padded[-1].item()) + output_tensor = torch.full( + (1, physical_tokens, 2), + float(bin_index + 1), + dtype=torch.bfloat16, + ) + wrapped = kwargs["post_processing_fn"]( + data_dict=processed_mb.data_dict, + cu_seqlens_padded=processed_mb.cu_seqlens_padded, + packed_seq_params=processed_mb.packed_seq_params, + ) + _, output = wrapped(output_tensor) + outputs.append(output) + retained_outputs.extend(outputs) + return outputs + + def fake_handle(storage): + return (tuple(storage.shape), int(storage.storage_offset())) + + monkeypatch.setattr(mpw, "get_microbatch_iterator", fake_get_microbatch_iterator) + monkeypatch.setattr( + mpw, "megatron_forward_backward", fake_megatron_forward_backward + ) + monkeypatch.setattr(mpw, "get_handle_from_tensor", fake_handle) + monkeypatch.setattr(mpw.torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(mpw.torch.cuda.nvtx, "range_push", lambda _name: None) + monkeypatch.setattr(mpw.torch.cuda.nvtx, "range_pop", lambda: None) + monkeypatch.setattr(mpw.torch.distributed, "get_rank", lambda: 7) + monkeypatch.setattr(mpw.parallel_state, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + mpw.parallel_state, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr(mpw.parallel_state, "get_data_parallel_rank", lambda: 0) + monkeypatch.setattr(mpw.parallel_state, "get_context_parallel_rank", lambda: 0) + monkeypatch.setattr( + mpw.parallel_state, "get_context_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr( + mpw.parallel_state, + "is_pipeline_last_stage", + lambda ignore_virtual: ignore_virtual, + ) + + worker = object.__new__(mpw.MegatronPolicyWorkerImpl) + worker.delegate_pack_to_model = False + worker.cfg = { + "logprob_batch_size": 1, + "sequence_packing": {"enabled": True}, + } + worker.model = SimpleNamespace(eval=lambda: None) + worker.mcore_state = SimpleNamespace(straggler_timer=None) + worker.defer_fp32_logits = False + worker.final_padded_vocab_size = 2 + flat_input_lengths = [value for microbatch in input_lengths for value in microbatch] + worker._teacher_ipc_storage = torch.empty(sum(flat_input_lengths), 2) + worker._teacher_ipc_handles = [] + + result = mpw.MegatronPolicyWorkerImpl.get_full_logits_ipc( + worker, + BatchedDataDict( + { + "input_ids": torch.zeros(sum(batch_sizes), 4, dtype=torch.long), + "input_lengths": torch.tensor(flat_input_lengths), + } + ), + micro_batch_size=1, + ) + + assert worker._teacher_ipc_storage.shape == (sum(flat_input_lengths), 2) + assert [record["batch_item_id"] for record in result["per_sample_handles"]] == [ + item_id for microbatch in batch_item_ids for item_id in microbatch + ] + assert [record["buf_idx"] for record in result["per_sample_handles"]] == [ + 0, + ] * sum(batch_sizes) + assert [ + record["sample_index_in_buf"] for record in result["per_sample_handles"] + ] == [index for batch_size in batch_sizes for index in range(batch_size)] + assert [ + record["storage_row_offset"] for record in result["per_sample_handles"] + ] == [0] * 7 + [7] * 2 + [9] + expected_token_offsets = [] + token_cursor = 0 + for valid_length in flat_input_lengths: + expected_token_offsets.append(token_cursor) + token_cursor += valid_length + assert [ + record["storage_token_offset"] for record in result["per_sample_handles"] + ] == expected_token_offsets + assert [ + record["stored_seq_len"] for record in result["per_sample_handles"] + ] == flat_input_lengths + assert [ + record["valid_seq_len"] for record in result["per_sample_handles"] + ] == flat_input_lengths + assert [record["actual_shape"] for record in result["per_sample_handles"]] == [ + (4, 2) + ] * sum(batch_sizes) + assert worker._teacher_ipc_handles == [((sum(flat_input_lengths), 2), 0)] + assert {record["payload_ipc"] for record in result["per_sample_handles"]} == { + worker._teacher_ipc_handles[0] + } + token_cursor = 0 + for expected_value, valid_length in zip( + [ + bin_index + 1 + for bin_index, batch_size in enumerate(batch_sizes) + for _ in range(batch_size) + ], + flat_input_lengths, + strict=True, + ): + assert torch.equal( + worker._teacher_ipc_storage[token_cursor : token_cursor + valid_length], + torch.full((valid_length, 2), float(expected_value)), + ) + token_cursor += valid_length + assert token_cursor == worker._teacher_ipc_storage.shape[0] + assert all( + isinstance(output["full_logits"], mpw.StreamedFullLogitsMetadata) + for output in retained_outputs + ) + assert all( + not isinstance(value, torch.Tensor) + for output in retained_outputs + for value in vars(output["full_logits"]).values() + ) + + +def test_full_logits_ipc_releases_and_grows_slab_before_forward(monkeypatch): + from megatron.core.packed_seq_params import PackedSeqParams + + from nemo_rl.models.policy.workers import megatron_policy_worker as mpw + + events = [] + input_lengths = [4, 3] + cu_seqlens = torch.tensor([0, 4, 7], dtype=torch.int32) + packed_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens, + cu_seqlens_kv_padded=cu_seqlens, + ) + processed_microbatch = SimpleNamespace( + data_dict=BatchedDataDict( + { + "batch_item_id": torch.tensor([10, 11]), + "input_ids": torch.zeros(2, 4, dtype=torch.long), + "input_lengths": torch.tensor(input_lengths), + } + ), + cu_seqlens_padded=cu_seqlens, + packed_seq_params=packed_params, + ) + + def fake_get_microbatch_iterator(*_args, **_kwargs): + return iter([processed_microbatch]), 1, 1, 4, 4 + + def fake_megatron_forward_backward(**kwargs): + processed_mb = next(kwargs["data_iterator"]) + events.append("forward") + assert worker._teacher_ipc_storage.shape == (sum(input_lengths), 2) + wrapped = kwargs["post_processing_fn"]( + data_dict=processed_mb.data_dict, + cu_seqlens_padded=processed_mb.cu_seqlens_padded, + packed_seq_params=processed_mb.packed_seq_params, + ) + _, output = wrapped(torch.ones(1, sum(input_lengths), 2)) + return [output] + + real_ensure = mpw.ensure_teacher_ipc_token_buffer + + def recording_ensure(storage, **kwargs): + events.append("allocate" if storage is None else "reuse") + if storage is None: + assert old_storage_ref() is None + assert worker._teacher_ipc_storage is None + assert worker._teacher_ipc_handles == [] + return real_ensure(storage, **kwargs) + + monkeypatch.setattr(mpw, "get_microbatch_iterator", fake_get_microbatch_iterator) + monkeypatch.setattr( + mpw, "megatron_forward_backward", fake_megatron_forward_backward + ) + monkeypatch.setattr(mpw, "ensure_teacher_ipc_token_buffer", recording_ensure) + monkeypatch.setattr( + mpw, "get_handle_from_tensor", lambda storage: (tuple(storage.shape), 0) + ) + monkeypatch.setattr(mpw.gc, "collect", lambda: events.append("collect")) + monkeypatch.setattr( + mpw.torch.cuda, "empty_cache", lambda: events.append("empty_cache") + ) + monkeypatch.setattr(mpw.torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(mpw.torch.cuda.nvtx, "range_push", lambda _name: None) + monkeypatch.setattr(mpw.torch.cuda.nvtx, "range_pop", lambda: None) + monkeypatch.setattr(mpw.torch.distributed, "get_rank", lambda: 7) + monkeypatch.setattr(mpw.parallel_state, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + mpw.parallel_state, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr(mpw.parallel_state, "get_data_parallel_rank", lambda: 0) + monkeypatch.setattr(mpw.parallel_state, "get_context_parallel_rank", lambda: 0) + monkeypatch.setattr( + mpw.parallel_state, "get_context_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr( + mpw.parallel_state, + "is_pipeline_last_stage", + lambda ignore_virtual: ignore_virtual, + ) + + worker = object.__new__(mpw.MegatronPolicyWorkerImpl) + worker.delegate_pack_to_model = False + worker.cfg = { + "logprob_batch_size": 1, + "sequence_packing": {"enabled": True}, + } + worker.model = SimpleNamespace(eval=lambda: None) + worker.mcore_state = SimpleNamespace(straggler_timer=None) + worker.defer_fp32_logits = False + worker.final_padded_vocab_size = 2 + old_storage = torch.empty(3, 2) + old_storage_ref = weakref.ref(old_storage) + worker._teacher_ipc_storage = old_storage + worker._teacher_ipc_handles = [("old",)] + del old_storage + + mpw.MegatronPolicyWorkerImpl.get_full_logits_ipc( + worker, + BatchedDataDict( + { + "input_ids": torch.zeros(2, 4, dtype=torch.long), + "input_lengths": torch.tensor(input_lengths), + } + ), + micro_batch_size=1, + ) + + assert events == ["collect", "empty_cache", "allocate", "forward"] + assert worker._teacher_ipc_storage.shape == (sum(input_lengths), 2) + + +def test_full_logits_ipc_nonlast_pipeline_stage_does_not_allocate(monkeypatch): + from nemo_rl.models.policy.workers import megatron_policy_worker as mpw + + processed_microbatch = SimpleNamespace( + data_dict=BatchedDataDict( + { + "batch_item_id": torch.tensor([10]), + "input_ids": torch.zeros(1, 4, dtype=torch.long), + "input_lengths": torch.tensor([3]), + } + ) + ) + + monkeypatch.setattr( + mpw, + "get_microbatch_iterator", + lambda *_args, **_kwargs: (iter([processed_microbatch]), 1, 1, 4, 4), + ) + + def fake_forward_backward(**kwargs): + list(kwargs["data_iterator"]) + assert kwargs["post_processing_fn"].packed_output_storage is None + return [] + + monkeypatch.setattr(mpw, "megatron_forward_backward", fake_forward_backward) + monkeypatch.setattr(mpw.torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(mpw.torch.cuda.nvtx, "range_push", lambda _name: None) + monkeypatch.setattr(mpw.torch.cuda.nvtx, "range_pop", lambda: None) + monkeypatch.setattr(mpw.torch.distributed, "get_rank", lambda: 7) + monkeypatch.setattr(mpw.parallel_state, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + mpw.parallel_state, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr(mpw.parallel_state, "get_data_parallel_rank", lambda: 0) + monkeypatch.setattr(mpw.parallel_state, "get_context_parallel_rank", lambda: 0) + monkeypatch.setattr( + mpw.parallel_state, "get_context_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr( + mpw.parallel_state, + "is_pipeline_last_stage", + lambda ignore_virtual: not ignore_virtual, + ) + + worker = object.__new__(mpw.MegatronPolicyWorkerImpl) + worker.delegate_pack_to_model = False + worker.cfg = { + "logprob_batch_size": 1, + "sequence_packing": {"enabled": True}, + } + worker.model = SimpleNamespace(eval=lambda: None) + worker.mcore_state = SimpleNamespace(straggler_timer=None) + worker.defer_fp32_logits = False + worker._teacher_ipc_storage = None + worker._teacher_ipc_handles = [] + worker._ensure_teacher_ipc_storage = MagicMock( + side_effect=AssertionError("non-last PP stage must not allocate IPC storage") + ) + + result = mpw.MegatronPolicyWorkerImpl.get_full_logits_ipc( + worker, + BatchedDataDict( + { + "input_ids": torch.zeros(1, 4, dtype=torch.long), + "input_lengths": torch.tensor([3]), + } + ), + micro_batch_size=1, + ) + + assert result == {"per_sample_handles": [], "dp_rank": 0} + worker._ensure_teacher_ipc_storage.assert_not_called() + + +def _sparse_ipc_worker(monkeypatch, *, tp_rank: int = 0, cp_rank: int = 1): + """A worker double whose forward returns fixed sparse rows per microbatch.""" + from nemo_rl.models.policy import utils as policy_utils + from nemo_rl.models.policy.workers import megatron_policy_worker as mpw + + batch_size, local_seq_len, k, force_width = 2, 4, 3, 2 + num_microbatches = 2 + + def fake_get_microbatch_iterator(*_args, **kwargs): + assert kwargs["skip_keys"] == {"force_include_token_ids"} + return iter([]), num_microbatches, batch_size, 8, 8 + + def fake_forward_backward(**kwargs): + assert kwargs["forward_only"] is True + assert kwargs["sampling_params"] is None + processor = kwargs["post_processing_fn"] + assert isinstance(processor, mpw.SparseTeacherLogitsPostProcessor) + outputs = [] + for buf_idx in range(num_microbatches): + base = float(10 * (buf_idx + 1) + len(calls)) + outputs.append( + { + "sparse_topk_logits": torch.full( + (batch_size, local_seq_len, k), base + ), + "sparse_topk_indices": torch.arange(k, dtype=torch.int32).expand( + batch_size, local_seq_len, k + ), + "sparse_log_z": torch.full((batch_size, local_seq_len), base + 1), + "sparse_natural_tail_indices": torch.full( + (batch_size, local_seq_len, 1), k - 1, dtype=torch.int32 + ), + "sparse_forced_logits": torch.full( + (batch_size, local_seq_len, force_width), base + 2 + ), + "sparse_forced_indices": torch.full( + (batch_size, local_seq_len, force_width), -1, dtype=torch.int32 + ), + "sparse_forced_in_topk": torch.zeros( + (batch_size, local_seq_len, force_width), dtype=torch.bool + ), + } + ) + calls.append(processor) + return outputs + + calls: list = [] + exported: list[int] = [] + + def fake_handle(storage): + exported.append(id(storage)) + return ("handle", id(storage)) + + monkeypatch.setattr(mpw, "get_microbatch_iterator", fake_get_microbatch_iterator) + monkeypatch.setattr(mpw, "megatron_forward_backward", fake_forward_backward) + monkeypatch.setattr(policy_utils, "get_handle_from_tensor", fake_handle) + monkeypatch.setattr(mpw.torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(mpw.torch.cuda.nvtx, "range_push", lambda _name: None) + monkeypatch.setattr(mpw.torch.cuda.nvtx, "range_pop", lambda: None) + monkeypatch.setattr(mpw.torch.distributed, "get_rank", lambda: 5) + monkeypatch.setattr( + mpw.parallel_state, "get_tensor_model_parallel_rank", lambda: tp_rank + ) + monkeypatch.setattr( + mpw.parallel_state, "get_tensor_model_parallel_world_size", lambda: 2 + ) + monkeypatch.setattr(mpw.parallel_state, "get_data_parallel_rank", lambda: 0) + monkeypatch.setattr( + mpw.parallel_state, "get_context_parallel_rank", lambda: cp_rank + ) + monkeypatch.setattr( + mpw.parallel_state, "get_context_parallel_world_size", lambda: 2 + ) + monkeypatch.setattr( + mpw.parallel_state, "is_pipeline_last_stage", lambda ignore_virtual: True + ) + + worker = object.__new__(mpw.MegatronPolicyWorkerImpl) + worker.delegate_pack_to_model = False + worker.delegate_mtp_loss_mask_to_model = False + worker.model_slices_context_parallel_inputs = False + worker.mtp_enabled = False + worker.cfg = { + "logprob_batch_size": batch_size, + "max_total_sequence_length": 12, + "sequence_packing": {"enabled": False}, + } + worker.model = SimpleNamespace(eval=lambda: None) + worker.mcore_state = SimpleNamespace(straggler_timer=None) + worker.defer_fp32_logits = False + worker.final_padded_vocab_size = 64 + worker._teacher_ipc_storage = None + worker._teacher_ipc_handles = [] + worker._teacher_sparse_ipc_storage = {} + worker._teacher_sparse_ipc_handles = {} + return mpw, worker, calls, exported + + +def _teacher_sparse_data() -> BatchedDataDict: + return BatchedDataDict( + { + "input_ids": torch.zeros((4, 8), dtype=torch.long), + "force_include_token_ids": torch.full((4, 8, 2), -1), + } + ) + + +def test_sparse_logits_ipc_writes_seven_fields_and_reuses_slabs(monkeypatch): + mpw, worker, calls, exported = _sparse_ipc_worker(monkeypatch) + + first = mpw.MegatronPolicyWorkerImpl.get_sparse_logits_ipc( + worker, + _teacher_sparse_data(), + k=3, + temperature=0.5, + real_vocab_size=50, + noise_filter_k=0, + ) + + processor = calls[0] + assert (processor.k, processor.temperature) == (3, 0.5) + assert (processor.real_vocab_size, processor.noise_filter_k) == (50, 0) + storage = worker._teacher_sparse_ipc_storage + assert set(storage) == { + "topk_logits", + "topk_indices", + "log_z", + "natural_tail_indices", + "forced_logits", + "forced_indices", + "forced_in_topk", + } + # Capacity is reserved for max_total_sequence_length=12 over CP=2. + assert storage["topk_logits"].shape == (2, 2, 6, 3) + assert storage["log_z"].shape == (2, 2, 6, 1) + assert storage["forced_indices"].shape == (2, 2, 6, 2) + assert storage["forced_in_topk"].dtype == torch.bool + assert torch.equal(storage["topk_logits"][1, :, :4], torch.full((2, 4, 3), 20.0)) + assert torch.equal(storage["log_z"][0, :, :4, 0], torch.full((2, 4), 11.0)) + assert len(exported) == 7 + + records = first["per_sample_handles"] + assert first["dp_rank"] == 0 + assert [(r["buf_idx"], r["sample_index_in_buf"]) for r in records] == [ + (0, 0), + (0, 1), + (1, 0), + (1, 1), + ] + record = records[0] + assert record["transport"] == "sparse_topk" + assert ( + record["topk_logits_ipc"] == worker._teacher_sparse_ipc_handles["topk_logits"] + ) + assert record["sequence_segments"] == [(0, 2, 2), (2, 4, 2)] + assert record["full_seq_len"] == 8 + assert record["topk_shape"] == (4, 3) + assert record["forced_shape"] == (4, 2) + assert record["scalar_shape"] == (4,) + assert (record["k"], record["force_width"], record["forced_in_topk_k"]) == ( + 3, + 2, + 3, + ) + assert (record["cp_rank"], record["cp_size"], record["tp_size"]) == (1, 2, 2) + assert record["padded_vocab_size"] == 64 + + second = mpw.MegatronPolicyWorkerImpl.get_sparse_logits_ipc( + worker, + _teacher_sparse_data(), + k=3, + temperature=0.5, + real_vocab_size=50, + noise_filter_k=0, + ) + assert worker._teacher_sparse_ipc_storage["topk_logits"] is storage["topk_logits"] + assert len(exported) == 7, "cached tickets must be reused, not re-exported" + assert torch.equal(storage["topk_logits"][0, :, :4], torch.full((2, 4, 3), 11.0)) + assert ( + second["per_sample_handles"][0]["log_z_ipc"] + == first["per_sample_handles"][0]["log_z_ipc"] + ) + + monkeypatch.setattr(mpw.gc, "collect", lambda: None) + monkeypatch.setattr(mpw.torch.cuda, "empty_cache", lambda: None) + mpw.MegatronPolicyWorkerImpl.release_ipc_buffer(worker) + assert worker._teacher_sparse_ipc_storage == {} + assert worker._teacher_sparse_ipc_handles == {} + + +def test_sparse_logits_ipc_only_tp_rank_zero_exports(monkeypatch): + mpw, worker, calls, exported = _sparse_ipc_worker(monkeypatch, tp_rank=1) + + result = mpw.MegatronPolicyWorkerImpl.get_sparse_logits_ipc( + worker, + _teacher_sparse_data(), + k=3, + temperature=1.0, + real_vocab_size=50, + noise_filter_k=4, + ) + + assert len(calls) == 1, "every TP rank joins the top-k collectives" + assert result == {"per_sample_handles": [], "dp_rank": 0} + assert worker._teacher_sparse_ipc_storage == {} + assert exported == [] + + +def test_sparse_logits_ipc_rejects_model_owned_cp_packing(monkeypatch): + mpw, worker, calls, _ = _sparse_ipc_worker(monkeypatch) + worker.delegate_pack_to_model = True + + with pytest.raises(NotImplementedError, match="context_parallel_size=1"): + mpw.MegatronPolicyWorkerImpl.get_sparse_logits_ipc( + worker, + _teacher_sparse_data(), + k=3, + temperature=1.0, + real_vocab_size=50, + noise_filter_k=0, + ) + assert calls == [] + + def test_megatron_prepare_for_training_restores_optimizer(): from nemo_rl.models.policy.workers.megatron_policy_worker import ( MegatronPolicyWorkerImpl, diff --git a/tests/unit/models/policy/test_policy_validation.py b/tests/unit/models/policy/test_policy_validation.py index 8a96d5e6053..7f8ff660429 100644 --- a/tests/unit/models/policy/test_policy_validation.py +++ b/tests/unit/models/policy/test_policy_validation.py @@ -24,7 +24,14 @@ from unittest.mock import MagicMock, patch import pytest +import torch +from nemo_rl.data.packing import ( + LockstepPackingItem, + SidePackingSpec, + build_lockstep_packing_plan, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.generation.vllm.patches import ( VLLM_NEMOTRON_H_FP32_LM_HEAD_ENV_VAR, ) @@ -89,6 +96,134 @@ def create_mock_tokenizer(): return tokenizer +@pytest.mark.parametrize("entrypoint", ["train", "full_logits"]) +def test_lockstep_plan_dp_must_match_policy_dp(entrypoint): + plan = build_lockstep_packing_plan( + batch_uid=5, + items=tuple( + LockstepPackingItem(sample_id=f"sample-{i}", batch_item_id=i) + for i in range(4) + ), + sides=( + SidePackingSpec( + side_id="student", + capacity=8, + raw_lengths=(2, 2, 2, 2), + effective_lengths=(2, 2, 2, 2), + ), + ), + data_parallel_size=2, + ) + data = BatchedDataDict( + { + "batch_item_id": torch.arange(4), + "input_ids": torch.zeros((4, 2), dtype=torch.long), + "input_lengths": torch.full((4,), 2, dtype=torch.long), + } + ) + policy = Policy.__new__(Policy) + policy.use_sequence_packing = True + policy.use_dynamic_batches = False + policy.sharding_annotations = MagicMock() + policy.sharding_annotations.get_axis_size.return_value = 1 + + with pytest.raises(ValueError, match=r"targets DP=2.*uses DP=1"): + if entrypoint == "train": + policy._shard_for_train( + data, + batch_size=4, + packing_plan=plan, + packing_side_id="student", + ) + else: + policy.get_full_logits_ipc( + data, + packing_plan=plan, + packing_side_id="student", + ) + + +def _sparse_ipc_policy(cfg: dict) -> Policy: + policy = Policy.__new__(Policy) + policy.cfg = cfg + policy.use_sequence_packing = False + policy.use_dynamic_batches = False + policy.sharding_annotations = MagicMock() + policy.sharding_annotations.get_axis_size.return_value = 2 + policy.worker_group = MagicMock() + return policy + + +def test_sparse_logits_ipc_routes_megatron_workers_and_keeps_cp_shards(): + policy = _sparse_ipc_policy({"megatron_cfg": {"enabled": True}}) + record = {"transport": "sparse_topk", "cp_rank": 0} + policy.worker_group.get_all_worker_results.return_value = [ + {"dp_rank": 0, "per_sample_handles": [record]}, + # TP rank 1 exports nothing; the last CP shard still contributes. + {"dp_rank": 0, "per_sample_handles": []}, + {"dp_rank": 0, "per_sample_handles": [{**record, "cp_rank": 1}]}, + {"dp_rank": 1, "per_sample_handles": [{**record, "cp_rank": 0}]}, + ] + data = BatchedDataDict( + { + "input_ids": torch.zeros((2, 4), dtype=torch.long), + "force_include_token_ids": torch.full((2, 4, 2), -1), + } + ) + + handles = policy.get_sparse_logits_ipc( + data, + k=128, + temperature=0.5, + real_vocab_size=1000, + noise_filter_k=0, + micro_batch_size=1, + ) + + call = policy.worker_group.run_all_workers_sharded_data.call_args + assert call.args == ("get_sparse_logits_ipc",) + assert call.kwargs["in_sharded_axes"] == ["data_parallel"] + assert call.kwargs["replicate_on_axes"] == [ + "context_parallel", + "tensor_parallel", + "pipeline_parallel", + ] + assert "output_is_replicated" not in call.kwargs + assert call.kwargs["common_kwargs"] == { + "k": 128, + "temperature": 0.5, + "real_vocab_size": 1000, + "noise_filter_k": 0, + "micro_batch_size": 1, + } + assert [len(entry["teacher_shards"]) for entry in handles] == [2, 1] + assert [shard["cp_rank"] for shard in handles[0]["teacher_shards"]] == [0, 1] + + +@pytest.mark.parametrize( + "cfg,packing,dynamic,match", + [ + ({"automodel_cfg": {"enabled": True}}, False, False, "only Megatron"), + ({"megatron_cfg": {"enabled": False}}, False, False, "only Megatron"), + ({"megatron_cfg": {"enabled": True}}, True, False, "sequence packing"), + ({"megatron_cfg": {"enabled": True}}, False, True, "dynamic batching"), + ], +) +def test_sparse_logits_ipc_rejects_unsupported_policies(cfg, packing, dynamic, match): + policy = _sparse_ipc_policy(cfg) + policy.use_sequence_packing = packing + policy.use_dynamic_batches = dynamic + with pytest.raises(NotImplementedError, match=match): + policy.get_sparse_logits_ipc( + BatchedDataDict({"input_ids": torch.zeros((2, 4), dtype=torch.long)}), + k=2, + temperature=1.0, + real_vocab_size=8, + noise_filter_k=0, + ) + policy.worker_group.run_all_workers_sharded_data.assert_not_called() + + def create_automodel_config( model_name: str, tp: int, pp: int = 1, cp: int = 1 ) -> PolicyConfig: diff --git a/tests/unit/models/policy/test_utils.py b/tests/unit/models/policy/test_utils.py index e591f220369..71d7211b84e 100644 --- a/tests/unit/models/policy/test_utils.py +++ b/tests/unit/models/policy/test_utils.py @@ -24,16 +24,27 @@ import torch import zmq +from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.policy.utils import ( + DENSE_TEACHER_IPC_FLAT_LAYOUT, DETERMINISTIC_ENV_VARS, IPCProtocol, aggregate_per_sample_handles, calculate_aligned_size, enable_deterministic_algorithms, ensure_teacher_ipc_buffer, + ensure_teacher_ipc_row_buffer, + ensure_teacher_ipc_token_buffer, + extract_batch_item_ids, + extract_teacher_ipc_valid_lengths, + get_dense_ipc_sequence_layout, get_megatron_checkpoint_dir, + localize_teacher_ipc_valid_lengths, + partition_teacher_ipc_row_buffer, + partition_teacher_ipc_token_buffer, rebuild_cuda_tensor_from_ipc, stream_weights_via_ipc_zmq_impl, + validate_compact_teacher_ipc_handle, ) @@ -585,6 +596,38 @@ def _run_stream_weights_roundtrip( class TestAggregatePerSampleHandles: + @staticmethod + def _dense_handle( + batch_item_id, + shard, + *, + tp_rank=0, + tp_size=1, + cp_rank=0, + cp_size=1, + full_seq_len=8, + full_vocab_size=12, + global_seq_start=None, + ): + local_seq_len = full_seq_len // cp_size + local_vocab_size = full_vocab_size // tp_size + if global_seq_start is None: + global_seq_start = cp_rank * local_seq_len + return { + "batch_item_id": batch_item_id, + "shard": shard, + "tp_rank": tp_rank, + "tp_size": tp_size, + "cp_rank": cp_rank, + "cp_size": cp_size, + "actual_shape": (local_seq_len, local_vocab_size), + "global_seq_start": global_seq_start, + "full_seq_len": full_seq_len, + "vocab_start_index": tp_rank * local_vocab_size, + "vocab_end_index": (tp_rank + 1) * local_vocab_size, + "full_vocab_size": full_vocab_size, + } + def test_orders_by_dp_rank(self): out = aggregate_per_sample_handles( [ @@ -615,6 +658,237 @@ def test_length_mismatch_raises(self): ] ) + def test_identity_aware_aggregation_uses_exact_canonical_order(self): + out = aggregate_per_sample_handles( + [ + { + "dp_rank": 1, + "per_sample_handles": [ + self._dense_handle(30, "a30", tp_rank=0, tp_size=2), + self._dense_handle(10, "a10", tp_rank=0, tp_size=2), + ], + }, + { + "dp_rank": 0, + "per_sample_handles": [ + self._dense_handle(20, "a20", tp_rank=0, tp_size=2) + ], + }, + { + "dp_rank": 1, + "per_sample_handles": [ + self._dense_handle(10, "b10", tp_rank=1, tp_size=2), + self._dense_handle(30, "b30", tp_rank=1, tp_size=2), + ], + }, + { + "dp_rank": 0, + "per_sample_handles": [ + self._dense_handle(20, "b20", tp_rank=1, tp_size=2) + ], + }, + ], + canonical_batch_item_ids=(10, 20, 30), + ) + + assert [item["batch_item_id"] for item in out] == [10, 20, 30] + assert [[h["shard"] for h in item["teacher_shards"]] for item in out] == [ + ["a10", "b10"], + ["a20", "b20"], + ["a30", "b30"], + ] + + @pytest.mark.parametrize( + "worker_results,match", + [ + ( + [{"dp_rank": 0, "per_sample_handles": [{"shard": "missing"}]}], + "carry batch_item_id", + ), + ( + [ + { + "dp_rank": 0, + "per_sample_handles": [ + {"batch_item_id": 1}, + {"batch_item_id": 1}, + ], + } + ], + "duplicate handles", + ), + ( + [ + { + "dp_rank": 0, + "per_sample_handles": [{"batch_item_id": 1}], + }, + { + "dp_rank": 0, + "per_sample_handles": [{"batch_item_id": 2}], + }, + ], + "inconsistent batch_item_id sets", + ), + ], + ) + def test_identity_aware_aggregation_rejects_invalid_records( + self, worker_results, match + ): + with pytest.raises(ValueError, match=match): + aggregate_per_sample_handles( + worker_results, + canonical_batch_item_ids=(1, 2), + ) + + def test_identity_aware_aggregation_rejects_missing_tp_cp_shard(self): + worker_results = [ + { + "dp_rank": 0, + "per_sample_handles": [ + self._dense_handle( + 1, + f"tp{tp_rank}cp{cp_rank}", + tp_rank=tp_rank, + tp_size=2, + cp_rank=cp_rank, + cp_size=2, + ) + ], + } + for tp_rank, cp_rank in ((0, 0), (0, 1), (1, 0)) + ] + with pytest.raises(ValueError, match="incomplete TP/CP shard coverage"): + aggregate_per_sample_handles(worker_results, canonical_batch_item_ids=(1,)) + + def test_identity_aware_aggregation_rejects_duplicate_coordinate(self): + duplicate = self._dense_handle(1, "duplicate") + with pytest.raises(ValueError, match="duplicate shard coordinate"): + aggregate_per_sample_handles( + [ + {"dp_rank": 0, "per_sample_handles": [duplicate]}, + {"dp_rank": 0, "per_sample_handles": [dict(duplicate)]}, + ], + canonical_batch_item_ids=(1,), + ) + + def test_identity_aware_aggregation_rejects_sequence_coverage_gap(self): + with pytest.raises(ValueError, match="gap or overlap"): + aggregate_per_sample_handles( + [ + { + "dp_rank": 0, + "per_sample_handles": [ + self._dense_handle(1, "cp0", cp_rank=0, cp_size=2) + ], + }, + { + "dp_rank": 0, + "per_sample_handles": [ + self._dense_handle( + 1, + "cp1", + cp_rank=1, + cp_size=2, + global_seq_start=3, + ) + ], + }, + ], + canonical_batch_item_ids=(1,), + ) + + def test_identity_aware_aggregation_rejects_vocab_coverage_gap(self): + tp0 = self._dense_handle(1, "tp0", tp_rank=0, tp_size=2) + tp1 = self._dense_handle(1, "tp1", tp_rank=1, tp_size=2) + tp1["vocab_start_index"] = 5 + tp1["vocab_end_index"] = 11 + with pytest.raises(ValueError, match="gap or overlap"): + aggregate_per_sample_handles( + [ + {"dp_rank": 0, "per_sample_handles": [tp0]}, + {"dp_rank": 0, "per_sample_handles": [tp1]}, + ], + canonical_batch_item_ids=(1,), + ) + + def test_identity_aware_aggregation_rejects_compact_valid_length_drift(self): + cp0 = self._dense_handle(1, "cp0", cp_rank=0, cp_size=2) + cp1 = self._dense_handle(1, "cp1", cp_rank=1, cp_size=2) + cp0.update( + { + "ipc_layout": DENSE_TEACHER_IPC_FLAT_LAYOUT, + "payload_ipc": "cp0", + "storage_shape": (4, 12), + "storage_token_offset": 0, + "storage_used_tokens": 4, + "storage_capacity_tokens": 4, + "stored_seq_len": 4, + "valid_seq_len": 5, + "dtype": torch.float32, + } + ) + cp1.update( + { + "ipc_layout": DENSE_TEACHER_IPC_FLAT_LAYOUT, + "payload_ipc": None, + "storage_shape": (0, 12), + "storage_token_offset": 0, + "storage_used_tokens": 0, + "storage_capacity_tokens": 0, + "stored_seq_len": 0, + "valid_seq_len": 4, + "dtype": torch.float32, + } + ) + + with pytest.raises(ValueError, match="disagree on valid_seq_len"): + aggregate_per_sample_handles( + [ + {"dp_rank": 0, "per_sample_handles": [cp0]}, + {"dp_rank": 0, "per_sample_handles": [cp1]}, + ], + canonical_batch_item_ids=(1,), + ) + + +class TestExtractBatchItemIds: + def test_accepts_tensor_and_validates_cardinality(self): + assert extract_batch_item_ids( + {"batch_item_id": torch.tensor([7, 3])}, 2, required=True + ) == [7, 3] + with pytest.raises(ValueError, match="cardinality"): + extract_batch_item_ids({"batch_item_id": [7]}, 2, required=True) + + def test_required_identity_is_not_synthesized(self): + assert extract_batch_item_ids({}, 2, required=False) is None + with pytest.raises(ValueError, match="requires batch_item_id"): + extract_batch_item_ids({}, 2, required=True) + + +class TestDenseIpcSequenceLayout: + def test_uses_logical_rectangle_not_packed_physical_bin(self): + data = BatchedDataDict( + { + "input_ids": torch.zeros(2, 6, dtype=torch.long), + "input_lengths": torch.tensor([3, 2]), + } + ) + # This is the physical scheduler geometry for a packed bin. It must not + # leak into dense teacher handle metadata, whose sequence axis is T_t=6. + data.micro_batch_lengths = [[12]] + + assert get_dense_ipc_sequence_layout(data, cp_rank=0, cp_size=2) == ( + 6, + 3, + 0, + ) + assert get_dense_ipc_sequence_layout(data, cp_rank=1, cp_size=2) == ( + 6, + 3, + 3, + ) + class TestEnsureTeacherIpcBuffer: def test_alloc_reuse_and_grow(self): @@ -626,6 +900,219 @@ def test_alloc_reuse_and_grow(self): s3, _ = ensure_teacher_ipc_buffer(s, h, 3, 1, 4, 8, torch.float32, dev) assert s3 is not s and s3.shape == (3, 1, 4, 8) + def test_max_cardinality_preallocation_keeps_handle_stable(self): + dev = torch.device("cpu") + storage, handle = ensure_teacher_ipc_buffer( + None, None, 3, 4, 6, 8, torch.float32, dev + ) + for variable_batch_size in (1, 4, 2): + next_storage, next_handle = ensure_teacher_ipc_buffer( + storage, + handle, + 3, + variable_batch_size, + 6, + 8, + torch.float32, + dev, + ) + assert next_storage is storage + assert next_handle is handle + + def test_capacity_seq_len_reserves_a_sequence_high_water_mark(self): + dev = torch.device("cpu") + storage, handle = ensure_teacher_ipc_buffer( + None, None, 2, 1, 4, 8, torch.float32, dev, capacity_seq_len=10 + ) + assert storage.shape == (2, 1, 10, 8) + for longer_seq_len in (5, 10): + next_storage, next_handle = ensure_teacher_ipc_buffer( + storage, + handle, + 2, + 1, + longer_seq_len, + 8, + torch.float32, + dev, + capacity_seq_len=10, + ) + assert next_storage is storage + assert next_handle is handle + grown, _ = ensure_teacher_ipc_buffer( + storage, handle, 2, 1, 12, 8, torch.float32, dev, capacity_seq_len=10 + ) + assert grown.shape == (2, 1, 12, 8) + + +class TestCompactTeacherIpcRowBuffer: + def test_variable_cardinality_bins_use_only_logical_rows(self): + # Qualification rank 3 is the worst old rectangular allocation: + # 6 bins * max(14 rows) = 84 rectangles for only 32 logical rows. + batch_sizes = [3, 5, 8, 14, 1, 1] + seq_len = 2 + vocab_size = 3 + storage = ensure_teacher_ipc_row_buffer( + None, + sum(batch_sizes), + seq_len, + vocab_size, + torch.float32, + torch.device("cpu"), + ) + partitions = partition_teacher_ipc_row_buffer(storage, batch_sizes) + + assert storage.shape == (32, seq_len, vocab_size) + assert len(batch_sizes) * max(batch_sizes) == 84 + assert storage.numel() == 32 * seq_len * vocab_size + + expected_offset = 0 + for bin_index, (row_offset, bin_view) in enumerate(partitions): + batch_size = batch_sizes[bin_index] + assert row_offset == expected_offset + assert bin_view.shape == (1, batch_size, seq_len, vocab_size) + assert bin_view.is_contiguous() + assert bin_view.storage_offset() == row_offset * seq_len * vocab_size + # This is the exact indexing ABI used by dense IPC consumers. + for sample_index_in_buf in range(batch_size): + bin_view[0, sample_index_in_buf].fill_(bin_index + 1) + assert torch.equal( + bin_view[0, sample_index_in_buf], + storage[row_offset + sample_index_in_buf], + ) + expected_offset += batch_size + assert expected_offset == storage.shape[0] + + def test_reuses_row_capacity_and_grows_before_partitioning(self): + device = torch.device("cpu") + storage = ensure_teacher_ipc_row_buffer(None, 8, 2, 3, torch.float32, device) + reused = ensure_teacher_ipc_row_buffer(storage, 5, 2, 3, torch.float32, device) + assert reused is storage + + grown = ensure_teacher_ipc_row_buffer(storage, 9, 2, 3, torch.float32, device) + assert grown is not storage + assert grown.shape == (9, 2, 3) + + with pytest.raises(ValueError, match="exceed storage capacity"): + partition_teacher_ipc_row_buffer(reused, [5, 4]) + + +class TestCompactTeacherIpcTokenBuffer: + def test_reuses_token_capacity_and_grows_only_when_required(self): + device = torch.device("cpu") + storage = ensure_teacher_ipc_token_buffer( + None, + total_tokens=8, + vocab_size=3, + dtype=torch.float32, + device=device, + ) + reused = ensure_teacher_ipc_token_buffer( + storage, + total_tokens=5, + vocab_size=3, + dtype=torch.float32, + device=device, + ) + grown = ensure_teacher_ipc_token_buffer( + storage, + total_tokens=9, + vocab_size=3, + dtype=torch.float32, + device=device, + ) + + assert reused is storage + assert grown is not storage + assert grown.shape == (9, 3) + + def test_qualification_like_capacity_is_sum_of_valid_cp_prefixes(self): + full_seq_len = 3912 + valid_lengths_by_microbatch = [ + [3912, 2744, 1512], + [2056, 1956, 1912, 264], + ] + flat_valid_lengths = [ + value for microbatch in valid_lengths_by_microbatch for value in microbatch + ] + localized_by_cp = [ + localize_teacher_ipc_valid_lengths( + flat_valid_lengths, + full_seq_len=full_seq_len, + cp_rank=cp_rank, + cp_size=2, + ) + for cp_rank in range(2) + ] + + assert sum(map(sum, localized_by_cp)) == sum(flat_valid_lengths) + for local_valid_lengths in localized_by_cp: + storage = ensure_teacher_ipc_token_buffer( + None, + total_tokens=sum(local_valid_lengths), + vocab_size=3, + dtype=torch.float32, + device=torch.device("cpu"), + ) + assert storage.shape == (sum(local_valid_lengths), 3) + assert storage.shape[0] < len(flat_valid_lengths) * (full_seq_len // 2) + + def test_partitions_unequal_rows_without_allocation_or_padding(self): + valid_lengths_by_microbatch = [[4, 2], [0, 3]] + storage = ensure_teacher_ipc_token_buffer( + None, + total_tokens=9, + vocab_size=2, + dtype=torch.float32, + device=torch.device("cpu"), + ) + partitions = partition_teacher_ipc_token_buffer( + storage, valid_lengths_by_microbatch + ) + + assert [[offset for offset, _ in values] for values in partitions] == [ + [0, 4], + [6, 6], + ] + assert [[tuple(view.shape) for _, view in values] for values in partitions] == [ + [(4, 2), (2, 2)], + [(0, 2), (3, 2)], + ] + assert all( + view.untyped_storage().data_ptr() == storage.untyped_storage().data_ptr() + for values in partitions + for _, view in values + ) + + def test_extracts_raw_lengths_only_for_packed_rows(self): + data = { + "input_lengths": torch.tensor([5, 2], dtype=torch.int64), + } + assert extract_teacher_ipc_valid_lengths( + data, 2, full_seq_len=8, compact_padding=True + ) == [5, 2] + assert extract_teacher_ipc_valid_lengths( + {}, 2, full_seq_len=8, compact_padding=False + ) == [8, 8] + + def test_compact_handle_rejects_out_of_bounds_offset(self): + handle = { + "ipc_layout": DENSE_TEACHER_IPC_FLAT_LAYOUT, + "payload_ipc": "payload", + "storage_shape": (4, 3), + "storage_token_offset": 3, + "storage_used_tokens": 4, + "storage_capacity_tokens": 4, + "stored_seq_len": 2, + "valid_seq_len": 2, + "actual_shape": (4, 3), + "global_seq_start": 0, + "full_seq_len": 4, + "dtype": torch.float32, + } + with pytest.raises(ValueError, match="outside storage bounds"): + validate_compact_teacher_ipc_handle(handle) + def test_enable_deterministic_algorithms_sets_torch_flags(): prev = ( diff --git a/tests/unit/test_config_validation.py b/tests/unit/test_config_validation.py index d92e47bcedc..ea9ce617f03 100644 --- a/tests/unit/test_config_validation.py +++ b/tests/unit/test_config_validation.py @@ -35,6 +35,9 @@ from nemo_rl.algorithms.ppo import MasterConfig as PPOMasterConfig from nemo_rl.algorithms.rm import MasterConfig as RMMasterConfig from nemo_rl.algorithms.sft import MasterConfig as SFTMasterConfig +from nemo_rl.algorithms.xtoken_off_policy_distillation import ( + CSDMasterConfig, +) from nemo_rl.algorithms.xtoken_off_policy_distillation import ( MasterConfig as XTokenOffPolicyDistillationMasterConfig, ) @@ -135,6 +138,9 @@ def test_all_config_files_have_required_keys(config_file): if "/evals/" in config_file: master_config_class = EvalMasterConfig config_type = "eval" + elif config_dict.get("algorithm") == "csd": + master_config_class = CSDMasterConfig + config_type = "csd" elif "teachers" in config_dict or ( "loss_fn" in config_dict and "projection_matrix_path" in config_dict["loss_fn"] ): diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py index fbfa1bd46b2..a4fcba83e85 100644 --- a/tests/unit/test_recipes_and_test_suites.py +++ b/tests/unit/test_recipes_and_test_suites.py @@ -43,6 +43,7 @@ # Relative to project root ALGO_MAPPING_TO_BASE_YAML = { + "csd": "examples/configs/csd.yaml", "sft": "examples/configs/sft.yaml", "dpo": "examples/configs/dpo.yaml", "grpo": "examples/configs/grpo_math_1B.yaml", @@ -286,7 +287,7 @@ def test_all_recipe_yamls_accounted_for_in_test_suites( ) -def test_nightly_compute_stays_below_4947_hours(nightly_test_suite, tracker): +def test_nightly_compute_stays_below_4953_hours(nightly_test_suite, tracker): command = f"DRYRUN=1 HF_HOME=... HF_DATASETS_CACHE=... CONTAINER= ACCOUNT= PARTITION= ./tools/launch {' '.join(nightly_test_suite)}" print(f"Running command: {command}") @@ -318,8 +319,8 @@ def test_nightly_compute_stays_below_4947_hours(nightly_test_suite, tracker): f"Last line of output was not as expected: '{last_line}'" ) total_gpu_hours = float(last_line.split(":")[-1].strip()) - assert total_gpu_hours <= 4947, ( - f"Total GPU hours exceeded 4947: {last_line}. We should revisit the test suites to reduce the total GPU hours." + assert total_gpu_hours <= 4953, ( + f"Total GPU hours exceeded 4953: {last_line}. We should revisit the test suites to reduce the total GPU hours." ) tracker.track("total_nightly_gpu_hours", total_gpu_hours) diff --git a/tests/unit/tools/csd/test_subtoks_tables.py b/tests/unit/tools/csd/test_subtoks_tables.py new file mode 100644 index 00000000000..dc385d8fa56 --- /dev/null +++ b/tests/unit/tools/csd/test_subtoks_tables.py @@ -0,0 +1,249 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU regression tests for CSD table preparation and validation.""" + +import sys +from pathlib import Path + +import pytest +import torch +from tokenizers import Tokenizer, decoders, models, pre_tokenizers +from transformers import PreTrainedTokenizerFast + +from tools.csd.build_subtoks_table import ( + build_subtoks_table, + load_tokenizer, + tokenizer_fingerprint, +) +from tools.csd.build_subtoks_table import ( + main as build_main, +) +from tools.csd.validate_subtoks_tables import main as validate_main +from tools.csd.validate_subtoks_tables import validate_table + + +def _byte_tokenizer( + extra_tokens: list[str], merges: list[tuple[str, str]] +) -> PreTrainedTokenizerFast: + surfaces = sorted(pre_tokenizers.ByteLevel.alphabet()) + extra_tokens + backend = Tokenizer( + models.BPE( + vocab={surface: i for i, surface in enumerate(surfaces)}, merges=merges + ) + ) + backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) + backend.decoder = decoders.ByteLevel() + return PreTrainedTokenizerFast(tokenizer_object=backend) + + +def _save_table( + path: Path, source: PreTrainedTokenizerFast, destination: PreTrainedTokenizerFast +) -> dict: + subtoks, lengths = build_subtoks_table(source, destination, max_subtokens=8) + table = { + "subtoks": subtoks, + "lengths": lengths, + "N_max": 8, + "V_src": len(source), + "V_dst": len(destination), + "source_fingerprint": tokenizer_fingerprint(source), + "destination_fingerprint": tokenizer_fingerprint(destination), + } + torch.save(table, path) + return table + + +def test_forward_and_reverse_tables_include_complete_chains(): + student = _byte_tokenizer(["ab"], [("a", "b")]) + teacher = _byte_tokenizer(["xy"], [("x", "y")]) + for source, destination, word in ( + (student, teacher, "ab"), + (teacher, student, "xy"), + ): + subtoks, lengths = build_subtoks_table(source, destination, max_subtokens=8) + source_id = source.convert_tokens_to_ids(word) + assert lengths[source_id].item() == 2 + assert subtoks[source_id, :2].tolist() == destination.convert_tokens_to_ids( + list(word) + ) + assert subtoks[source_id, 2:].tolist() == [-1] * 6 + + +def test_partial_utf8_is_preserved_as_bytes(): + source = _byte_tokenizer(["âĢ"], [("â", "Ģ")]) + destination = _byte_tokenizer(["ï¿", "�"], [("ï", "¿"), ("ï¿", "½")]) + source_id = source.convert_tokens_to_ids("âĢ") + assert source.decode([source_id]) == "\ufffd" + subtoks, lengths = build_subtoks_table(source, destination, max_subtokens=8) + assert lengths[source_id].item() == 2 + assert subtoks[source_id, :2].tolist() == destination.convert_tokens_to_ids( + ["â", "Ģ"] + ) + + +def test_overflow_is_unmapped_instead_of_truncating_the_event(): + source = _byte_tokenizer(["ab"], [("a", "b")]) + destination = _byte_tokenizer([], []) + subtoks, lengths = build_subtoks_table(source, destination, max_subtokens=1) + source_id = source.convert_tokens_to_ids("ab") + assert lengths[source_id].item() == 0 + assert subtoks[source_id].tolist() == [-1] + + +def test_nonpositive_chain_limit_is_rejected(): + tokenizer = _byte_tokenizer([], []) + with pytest.raises(ValueError, match="must be positive"): + build_subtoks_table(tokenizer, tokenizer, max_subtokens=0) + + +@pytest.mark.parametrize("unsupported_side", ["source", "destination"]) +@pytest.mark.parametrize("decoder_family", ["metaspace", "wordpiece"]) +def test_context_dependent_decoders_are_rejected_before_table_io( + tmp_path: Path, unsupported_side: str, decoder_family: str +): + if decoder_family == "metaspace": + backend = Tokenizer( + models.Unigram([("x", 0.0), ("▁the", -1.0), ("the", -2.0)], unk_id=0) + ) + backend.decoder = decoders.Metaspace(replacement="▁", prepend_scheme="always") + else: + backend = Tokenizer(models.WordPiece({"[UNK]": 0, "x": 1, "##the": 2})) + backend.decoder = decoders.WordPiece(prefix="##") + unsupported = PreTrainedTokenizerFast(tokenizer_object=backend) + if decoder_family == "metaspace": + # Isolated decoding loses the leading space, so a decoded-text-only + # table would falsely equate this token with the byte-BPE token "the". + assert unsupported.decode([1]) == "the" + assert unsupported.decode([0, 1]) == "x the" + byte_tokenizer = _byte_tokenizer([], []) + source, destination = ( + (unsupported, byte_tokenizer) + if unsupported_side == "source" + else (byte_tokenizer, unsupported) + ) + error = "BPE model and ByteLevel decoder" + with pytest.raises(ValueError, match=error): + build_subtoks_table(source, destination, max_subtokens=8) + with pytest.raises(ValueError, match=error): + validate_table(tmp_path / "table.pt", source=source, destination=destination) + + snapshot = tmp_path / "unsupported_tokenizer" + unsupported.save_pretrained(snapshot) + with pytest.raises(ValueError, match=error): + load_tokenizer(str(snapshot), revision=None) + + +def test_semantic_roles_precede_shared_spelling_and_missing_roles_are_unmapped(): + source = _byte_tokenizer([], []) + source.add_special_tokens({"eos_token": "", "bos_token": ""}) + destination = _byte_tokenizer([], []) + destination.add_special_tokens({"pad_token": "", "eos_token": ""}) + subtoks, lengths = build_subtoks_table(source, destination, max_subtokens=8) + assert lengths[source.eos_token_id].item() == 1 + assert subtoks[source.eos_token_id, 0].item() == destination.eos_token_id + assert lengths[source.bos_token_id].item() == 0 + + +def test_ambiguous_role_alias_uses_compatible_exact_special_spelling(): + source = _byte_tokenizer([], []) + source.add_special_tokens({"eos_token": "", "pad_token": ""}) + destination = _byte_tokenizer([], []) + destination.add_special_tokens({"eos_token": "", "pad_token": ""}) + subtoks, lengths = build_subtoks_table(source, destination, max_subtokens=8) + assert lengths[source.eos_token_id].item() == 1 + assert subtoks[source.eos_token_id, 0].item() == destination.eos_token_id + + +def test_ordinary_added_token_cannot_become_a_semantic_special(): + source = _byte_tokenizer([], []) + source.add_tokens([""]) + destination = _byte_tokenizer([], []) + destination.add_special_tokens({"eos_token": ""}) + _, lengths = build_subtoks_table(source, destination, max_subtokens=8) + assert lengths[source.convert_tokens_to_ids("")].item() == 0 + + +@pytest.mark.parametrize( + "corruption", ["event", "range", "padding", "fingerprint", "role"] +) +def test_validator_rejects_invalid_tables(tmp_path: Path, corruption: str): + source = _byte_tokenizer(["ab"], [("a", "b")]) + destination = _byte_tokenizer([], []) + source.add_special_tokens({"eos_token": ""}) + destination.add_special_tokens({"eos_token": ""}) + path = tmp_path / "table.pt" + table = _save_table(path, source, destination) + validate_table(path, source=source, destination=destination) + row = source.convert_tokens_to_ids("ab") + if corruption == "event": + table["subtoks"][row, 1] = destination.convert_tokens_to_ids("x") + elif corruption == "range": + table["subtoks"][row, 0] = len(destination) + elif corruption == "padding": + table["subtoks"][row, 7] = 0 + elif corruption == "fingerprint": + table["source_fingerprint"] = "wrong-tokenizer" + elif corruption == "role": + table["subtoks"][source.eos_token_id, 0] = 0 + torch.save(table, path) + with pytest.raises(ValueError): + validate_table(path, source=source, destination=destination) + + +def test_public_build_and_validate_cli_workflow(tmp_path: Path, monkeypatch): + student = _byte_tokenizer(["ab"], [("a", "b")]) + teacher = _byte_tokenizer(["xy"], [("x", "y")]) + student_dir, teacher_dir = tmp_path / "student", tmp_path / "teacher" + student.save_pretrained(student_dir) + teacher.save_pretrained(teacher_dir) + forward, reverse = ( + tmp_path / "student_to_teacher.pt", + tmp_path / "teacher_to_student.pt", + ) + for source, destination, output in ( + (student_dir, teacher_dir, forward), + (teacher_dir, student_dir, reverse), + ): + monkeypatch.setattr( + sys, + "argv", + [ + "build_subtoks_table", + "--source-tokenizer", + str(source), + "--destination-tokenizer", + str(destination), + "--max-subtokens", + "8", + "--output", + str(output), + ], + ) + build_main() + monkeypatch.setattr( + sys, + "argv", + [ + "validate_subtoks_tables", + "--student-tokenizer", + str(student_dir), + "--teacher-tokenizer", + str(teacher_dir), + "--student-to-teacher", + str(forward), + "--teacher-to-student", + str(reverse), + ], + ) + validate_main() diff --git a/tools/csd/__init__.py b/tools/csd/__init__.py new file mode 100644 index 00000000000..2a616c2b1b1 --- /dev/null +++ b/tools/csd/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Offline preparation of CSD tokenizer sub-token tables.""" diff --git a/tools/csd/build_subtoks_table.py b/tools/csd/build_subtoks_table.py new file mode 100644 index 00000000000..c69d2fe6954 --- /dev/null +++ b/tools/csd/build_subtoks_table.py @@ -0,0 +1,303 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Build one direction of the CSD sub-token tables on CPU. + +Invoke once for student-to-teacher and once with the tokenizers reversed. +Each source vocabulary row contains a complete destination token chain; +unmapped rows have length zero and padding is -1. No model weights are loaded. +Both tokenizers must use a BPE model with a ByteLevel decoder. +""" + +import argparse +import hashlib +import json +from pathlib import Path + +import torch +from tokenizers.decoders import ByteLevel +from tokenizers.models import BPE +from transformers import AutoTokenizer, PreTrainedTokenizerFast + + +def special_roles(tokenizer: PreTrainedTokenizerFast) -> dict[str, int | None]: + """Return semantic roles independently of their token spellings.""" + return { + "bos": tokenizer.bos_token_id, + "eos": tokenizer.eos_token_id, + "pad": tokenizer.pad_token_id, + "unk": tokenizer.unk_token_id, + "sep": tokenizer.sep_token_id, + "cls": tokenizer.cls_token_id, + "mask": tokenizer.mask_token_id, + } + + +def hard_special_ids(tokenizer: PreTrainedTokenizerFast) -> set[int]: + """Include hard added specials without excluding ordinary added tokens.""" + return set(tokenizer.all_special_ids) | { + token_id + for token_id, token in tokenizer.added_tokens_decoder.items() + if token.special + } + + +def tokenizer_fingerprint(tokenizer: PreTrainedTokenizerFast) -> str: + """Fingerprint vocabulary, tokenizer behavior, and semantic special roles.""" + backend = json.loads(tokenizer.backend_tokenizer.to_str()) + # Encode calls update these transient batch options; neither changes the + # vocabulary events represented by a table. + backend.pop("padding", None) + backend.pop("truncation", None) + payload = { + "backend": backend, + "roles": special_roles(tokenizer), + } + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest() + + +def byte_bpe_model(tokenizer: PreTrainedTokenizerFast) -> BPE | None: + """Return the byte-preserving BPE model, when this tokenizer has one.""" + backend = tokenizer.backend_tokenizer + if isinstance(backend.model, BPE) and isinstance(backend.decoder, ByteLevel): + return backend.model + return None + + +def require_supported_tokenizer( + tokenizer: PreTrainedTokenizerFast, *, label: str +) -> None: + """Reject decoders whose single-token text can change with context. + + Metaspace and WordPiece decoders can add or remove whitespace or word + markers depending on token position. Isolated decoding cannot establish + the complete events needed by a static table for those tokenizers. + """ + if byte_bpe_model(tokenizer) is None: + raise ValueError( + f"{label}: CSD table preparation and validation require a fast " + "tokenizer with a BPE model and ByteLevel decoder. Other decoder " + "families, including Metaspace and WordPiece, are unsupported " + "because isolated token decoding can change events with context." + ) + + +def role_mapping( + source: PreTrainedTokenizerFast, destination: PreTrainedTokenizerFast +) -> dict[int, int]: + """Map unambiguous shared special roles; never resolve aliases by order.""" + destination_roles = special_roles(destination) + candidates: dict[int, set[int]] = {} + for role, source_id in special_roles(source).items(): + destination_id = destination_roles[role] + if source_id is not None and destination_id is not None: + candidates.setdefault(source_id, set()).add(destination_id) + return { + source_id: next(iter(targets)) + for source_id, targets in candidates.items() + if len(targets) == 1 + } + + +def special_match_is_valid( + source: PreTrainedTokenizerFast, + destination: PreTrainedTokenizerFast, + *, + source_id: int, + destination_id: int, +) -> bool: + """Permit identical hard-special spellings only when roles do not conflict.""" + source_roles = {k for k, v in special_roles(source).items() if v == source_id} + destination_roles = { + k for k, v in special_roles(destination).items() if v == destination_id + } + return not ( + source_roles + and destination_roles + and source_roles.isdisjoint(destination_roles) + ) and source.convert_ids_to_tokens(source_id) == destination.convert_ids_to_tokens( + destination_id + ) + + +def ordinary_event_matches( + source: PreTrainedTokenizerFast, + destination: PreTrainedTokenizerFast, + *, + source_id: int, + destination_ids: list[int], +) -> bool: + """Check complete events without replacing incomplete UTF-8 fragments. + + Byte-level BPE pairs compare their reversible byte surfaces. Ordinary added + tokens compare decoded text and conservatively reject replacement characters. + """ + if not destination_ids: + return False + if ( + byte_bpe_model(source) is not None + and byte_bpe_model(destination) is not None + and source_id not in source.added_tokens_decoder + and not any(i in destination.added_tokens_decoder for i in destination_ids) + ): + return source.convert_ids_to_tokens(source_id) == "".join( + destination.convert_ids_to_tokens(destination_ids) + ) + source_text = source.decode([source_id], clean_up_tokenization_spaces=False) + destination_text = destination.decode( + destination_ids, clean_up_tokenization_spaces=False + ) + return ( + bool(source_text) + and "\ufffd" not in source_text + and source_text == destination_text + ) + + +def build_subtoks_table( + source: PreTrainedTokenizerFast, + destination: PreTrainedTokenizerFast, + *, + max_subtokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build complete destination chains for every source vocabulary ID. + + Args: + source: Source tokenizer, including its added tokens. + destination: Tokenizer whose IDs populate each chain. + max_subtokens: Positive maximum chain length; longer events are unmapped. + + Returns: + Integer tensors ``subtoks[V_source, max_subtokens]`` and + ``lengths[V_source]``, ready for the CSD loss table loader. + """ + if max_subtokens < 1: + raise ValueError("max_subtokens must be positive") + require_supported_tokenizer(source, label="source tokenizer") + require_supported_tokenizer(destination, label="destination tokenizer") + subtoks = torch.full((len(source), max_subtokens), -1, dtype=torch.long) + lengths = torch.zeros(len(source), dtype=torch.long) + source_specials = hard_special_ids(source) + destination_specials = hard_special_ids(destination) + roles = role_mapping(source, destination) + destination_vocab = destination.get_vocab() + destination_bpe = byte_bpe_model(destination) + byte_pair = byte_bpe_model(source) is not None and destination_bpe is not None + + for source_id in range(len(source)): + surface = source.convert_ids_to_tokens(source_id) + exact_id = destination_vocab.get(surface) + destination_ids: list[int] = [] + if source_id in source_specials: + if source_id in roles: + destination_ids = [roles[source_id]] + elif ( + exact_id is not None + and exact_id in destination_specials + and special_match_is_valid( + source, destination, source_id=source_id, destination_id=exact_id + ) + ): + destination_ids = [exact_id] + else: + if ( + exact_id is not None + and exact_id not in destination_specials + and ordinary_event_matches( + source, destination, source_id=source_id, destination_ids=[exact_id] + ) + ): + destination_ids = [exact_id] + else: + text = source.decode([source_id], clean_up_tokenization_spaces=False) + candidate = destination.encode(text, add_special_tokens=False) + if ( + byte_pair + and source_id not in source.added_tokens_decoder + and not ordinary_event_matches( + source, + destination, + source_id=source_id, + destination_ids=candidate, + ) + ): + # Token-level Unicode decoding can replace incomplete byte + # fragments. BPE on the reversible surface preserves them. + assert destination_bpe is not None + candidate = [ + piece.id for piece in destination_bpe.tokenize(surface) + ] + if not any( + i in destination_specials for i in candidate + ) and ordinary_event_matches( + source, destination, source_id=source_id, destination_ids=candidate + ): + destination_ids = candidate + if 0 < len(destination_ids) <= max_subtokens: + lengths[source_id] = len(destination_ids) + subtoks[source_id, : len(destination_ids)] = torch.tensor(destination_ids) + return subtoks, lengths + + +def load_tokenizer(name: str, *, revision: str | None) -> PreTrainedTokenizerFast: + """Load only tokenizer assets and reject unsupported slow implementations.""" + tokenizer = AutoTokenizer.from_pretrained(name, revision=revision, use_fast=True) + if not isinstance(tokenizer, PreTrainedTokenizerFast): + raise ValueError("CSD table preparation requires a fast tokenizer") + require_supported_tokenizer(tokenizer, label=name) + return tokenizer + + +def main() -> None: + """Build and save a directional table with tokenizer identity metadata.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source-tokenizer", required=True) + parser.add_argument("--destination-tokenizer", required=True) + parser.add_argument("--source-revision", default=None) + parser.add_argument("--destination-revision", default=None) + parser.add_argument("--max-subtokens", type=int, default=8) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + if args.output.exists() and not args.overwrite: + parser.error(f"{args.output} already exists; use --overwrite to replace it") + source = load_tokenizer(args.source_tokenizer, revision=args.source_revision) + destination = load_tokenizer( + args.destination_tokenizer, revision=args.destination_revision + ) + subtoks, lengths = build_subtoks_table( + source, destination, max_subtokens=args.max_subtokens + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "subtoks": subtoks, + "lengths": lengths, + "src": args.source_tokenizer, + "dst": args.destination_tokenizer, + "N_max": args.max_subtokens, + "V_src": len(source), + "V_dst": len(destination), + "source_fingerprint": tokenizer_fingerprint(source), + "destination_fingerprint": tokenizer_fingerprint(destination), + }, + args.output, + ) + print( + f"Saved {args.output}: {int((lengths > 0).sum())}/{len(source)} mapped rows, " + f"{int((lengths == 0).sum())} unmapped rows" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/csd/validate_subtoks_tables.py b/tools/csd/validate_subtoks_tables.py new file mode 100644 index 00000000000..334d80a3e82 --- /dev/null +++ b/tools/csd/validate_subtoks_tables.py @@ -0,0 +1,144 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Validate both CSD tables against the exact student and teacher tokenizers.""" + +import argparse +from pathlib import Path + +import torch +from transformers import PreTrainedTokenizerFast + +from tools.csd.build_subtoks_table import ( + hard_special_ids, + load_tokenizer, + ordinary_event_matches, + require_supported_tokenizer, + role_mapping, + special_match_is_valid, + tokenizer_fingerprint, +) + + +def validate_table( + path: Path, + *, + source: PreTrainedTokenizerFast, + destination: PreTrainedTokenizerFast, +) -> None: + """Check tensor structure, tokenizer identity, roles, and every mapped event. + + Args: + path: Table saved by ``tools.csd.build_subtoks_table``. + source: Tokenizer whose vocabulary indexes rows. + destination: Tokenizer whose IDs populate each chain. + + Raises: + ValueError: The table is inconsistent or maps a different token event. + """ + require_supported_tokenizer(source, label="source tokenizer") + require_supported_tokenizer(destination, label="destination tokenizer") + table = torch.load(path, map_location="cpu", weights_only=True) + if not isinstance(table, dict): + raise ValueError(f"{path}: expected a table dictionary") + for key, expected in ( + ("V_src", len(source)), + ("V_dst", len(destination)), + ("source_fingerprint", tokenizer_fingerprint(source)), + ("destination_fingerprint", tokenizer_fingerprint(destination)), + ): + if table.get(key) != expected: + raise ValueError( + f"{path}: {key} does not match the supplied tokenizer pair" + ) + subtoks, lengths = table.get("subtoks"), table.get("lengths") + if not isinstance(subtoks, torch.Tensor) or not isinstance(lengths, torch.Tensor): + raise ValueError(f"{path}: missing subtoks or lengths tensors") + if subtoks.dtype != torch.long or lengths.dtype != torch.long: + raise ValueError(f"{path}: subtoks and lengths must be int64") + if ( + subtoks.ndim != 2 + or subtoks.shape[0] != len(source) + or subtoks.shape[1] < 1 + or lengths.shape != (len(source),) + or table.get("N_max") != subtoks.shape[1] + ): + raise ValueError(f"{path}: invalid tensor shapes or N_max metadata") + if bool(((lengths < 0) | (lengths > subtoks.shape[1])).any()): + raise ValueError(f"{path}: invalid chain lengths") + valid = torch.arange(subtoks.shape[1])[None, :] < lengths[:, None] + if bool(((subtoks[valid] < 0) | (subtoks[valid] >= len(destination))).any()): + raise ValueError(f"{path}: destination token ID out of range") + if bool((subtoks[~valid] != -1).any()): + raise ValueError(f"{path}: padding must be -1") + + roles = role_mapping(source, destination) + source_specials = hard_special_ids(source) + destination_specials = hard_special_ids(destination) + for source_id, destination_id in roles.items(): + if ( + lengths[source_id].item() != 1 + or subtoks[source_id, 0].item() != destination_id + ): + raise ValueError( + f"{path}: source row {source_id} has an incorrect semantic role" + ) + for source_id in (lengths > 0).nonzero().flatten().tolist(): + destination_ids = subtoks[source_id, : lengths[source_id]].tolist() + if source_id in source_specials: + valid_event = len(destination_ids) == 1 and ( + roles.get(source_id) == destination_ids[0] + or ( + destination_ids[0] in destination_specials + and special_match_is_valid( + source, + destination, + source_id=source_id, + destination_id=destination_ids[0], + ) + ) + ) + else: + valid_event = not any( + i in destination_specials for i in destination_ids + ) and ordinary_event_matches( + source, + destination, + source_id=source_id, + destination_ids=destination_ids, + ) + if not valid_event: + raise ValueError( + f"{path}: source row {source_id} does not preserve its token event" + ) + print(f"Validated {path}: {int((lengths > 0).sum())} mapped rows") + + +def main() -> None: + """Validate directional tables before starting a GPU training job.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--student-tokenizer", required=True) + parser.add_argument("--teacher-tokenizer", required=True) + parser.add_argument("--student-revision", default=None) + parser.add_argument("--teacher-revision", default=None) + parser.add_argument("--student-to-teacher", type=Path, required=True) + parser.add_argument("--teacher-to-student", type=Path, required=True) + args = parser.parse_args() + student = load_tokenizer(args.student_tokenizer, revision=args.student_revision) + teacher = load_tokenizer(args.teacher_tokenizer, revision=args.teacher_revision) + validate_table(args.student_to_teacher, source=student, destination=teacher) + validate_table(args.teacher_to_student, source=teacher, destination=student) + + +if __name__ == "__main__": + main() diff --git a/tools/x_token/build_projection_matrix.sh b/tools/x_token/build_projection_matrix.sh index f31a06ee413..8cafeac449a 100755 --- a/tools/x_token/build_projection_matrix.sh +++ b/tools/x_token/build_projection_matrix.sh @@ -15,7 +15,7 @@ # Single entrypoint that chains the three projection-prep CLIs under # tools/x_token/ to produce a runtime projection matrix from a -# (student, teacher) tokenizer pair. See docs/guides/xtoken-off-policy-distillation.md +# (student, teacher) tokenizer pair. See docs/guides/cross-tokenizer-distillation.md # for the underlying steps. set -euo pipefail