feat(draft): support packed context-parallel co-training

docs/design-docs/dflash-draft-body.md

@@ -0,0 +1,39 @@
+# DFlash Draft Body
+
+The target model provides hidden states for an input sequence. The DFlash draft
+body uses these states as context for short blocks of draft-token positions.
+It updates all positions in a block together and returns hidden states for a
+separate draft-loss calculation. The body does not produce vocabulary logits.
+
+## How It Works
+
+1. **Choose anchors.** An anchor is a valid position in the target sequence.
+   `build_dflash_batch_plan` chooses anchors and records which draft positions
+   belong to each one. The same input mask, sample IDs, optimizer step, and seed
+   produce the same plan. Rows with the same sample ID get the same anchors, so
+   give each rollout row a unique sample ID. A sample with no valid window
+   contributes no draft loss.
+2. **Build draft blocks.** Each block has one anchor position followed by
+   `gamma` draft positions. The caller supplies an embedding for each position.
+   The target model supplies selected hidden states from its layers.
+3. **Update the blocks.** `DFlashBody` projects the target hidden states and
+   passes the block embeddings through draft decoder layers. Within a layer,
+   each block can read the target prefix before its anchor and all positions
+   in that same block. It cannot read another block. Attention computes outputs
+   for block positions only. It does not update the target sequence.
+4. **Compute the loss.** The body returns one hidden state per block position.
+   The caller uses `plan.loss_mask` to exclude the anchor and invalid positions
+   from the draft loss. The body returns zero for invalid positions.
+
+With this separation, the caller can use the same target-model context for
+all draft blocks. The attention of each block stays independent.
+
+## Scope
+
+The body implements the draft-state update only. The caller supplies token
+embeddings and uses the returned states to compute vocabulary logits and loss.
+
+The body supports tensor parallelism. The tensor-parallel size must divide
+`num_key_value_heads`, so the default config supports a maximum of 8. The body
+does not support sequence parallelism or context parallelism greater than 1.
+The caller controls pipeline-parallel placement.

docs/design-docs/dspark-objectives.md

@@ -0,0 +1,24 @@
+# DSpark Objectives
+
+`dspark_tiled_objective` computes three losses for valid draft slots:
+
+| Loss | Purpose |
+| --- | --- |
+| Hard cross-entropy (CE) | Train on sampled draft labels. |
+| Total variation (TV) | Compare the draft and target token distributions. |
+| Confidence | Predict token acceptance when a confidence head is present. |
+
+The function processes draft tokens in bounded chunks. It keeps target
+logits and the target output head stop-gradient. The Markov head uses
+previous-token IDs from the target vocabulary. Loss labels use the selected
+draft-vocabulary order. The two vocabularies can have different sizes.
+
+Each loss returns raw numerator and valid-slot count bins. The combined loss
+uses the configured CE, TV, and confidence weights. Position weights scale
+both numerators and denominators. Reduce counts across data-parallel ranks
+before calling `DSparkLossBins.normalized`.
+
+`build_dspark_provider` attaches the Markov head and optional confidence head
+to a checkpointable draft body. PR #3711 defines the head shapes and
+checkpoint keys.
+Training-loop configuration and rollout execution are outside this change.

docs/guides/dflash-dspark-speculative-decoding.md

@@ -0,0 +1,213 @@
+# Train DFlash and DSpark Online
+
+DFlash and DSpark are block-parallel speculative decoders. NeMo RL trains either
+draft model with a Megatron policy. It then refits both models into vLLM for
+rollout generation.
+
+This guide covers the supported online co-training path. For GRPO setup, see
+the [GRPO guide](grpo.md). For the Eagle3 workflow, which has different packing
+constraints, see [Train with Eagle3 Speculative Decoding](eagle3-speculative-decoding.md).
+
+## Supported Topologies
+
+Use one of the layouts below. NeMo RL rejects unsupported layouts before
+training.
+
+| Component | Supported configuration |
+| --- | --- |
+| Training backend | Megatron Core |
+| Target tensor parallelism | TP2 |
+| Target pipeline parallelism | PP1 with no virtual pipeline parallelism |
+| Target context parallelism | CP1, CP2, or CP4 |
+| Sequence packing | Optional at CP1; required at CP2 and CP4 |
+| Target sequence parallelism | Supported only with sequence packing |
+| DFlash/DSpark body sequence parallelism | Disabled |
+| Generation topology | CP1 and PP1 |
+
+Set `use_fused_linear_logprobs: false`. For CP2 or CP4, use the split
+begin/microbatch/finish training path. Run `examples/run_grpo.py` with
+`data_plane.enabled=true` to select the synchronous TransferQueue trainer,
+which uses that path.
+
+## Configure DFlash
+
+The repository includes a one-node Qwen3-8B recipe at
+`examples/configs/recipes/llm/grpo-qwen3-8b-1n8g-megatron-dflash.yaml`.
+The core DFlash configuration is:
+
+```yaml
+policy:
+  model_name: Qwen/Qwen3-8B
+  megatron_cfg:
+    enabled: true
+    tensor_model_parallel_size: 2
+    pipeline_model_parallel_size: 1
+    context_parallel_size: 1
+    sequence_parallel: false
+    use_fused_linear_logprobs: false
+  dtensor_cfg:
+    enabled: false
+  sequence_packing:
+    enabled: false
+  draft:
+    speculator_type: dflash
+    enabled: true
+    model_name: z-lab/Qwen3-8B-DFlash-b16
+    loss_weight: 1.0
+    gamma: 5
+    anchors_per_sample: 2
+    mask_token_id: 151669
+    target_hidden_state_layer_ids: [1, 9, 17, 25, 33]
+    num_layers: 5
+    seed: 13
+    vocab_tile_size: 256
+    position_decay: 1.0
+    max_cp_boundary_exclusion_fraction: 0.25
+  generation:
+    backend: vllm
+    vllm_kwargs:
+      speculative_config:
+        method: dflash
+        model: ${policy.draft.model_name}
+        num_speculative_tokens: ${policy.draft.gamma}
+        draft_tensor_parallel_size: 1
+```
+
+`gamma` is the number of masked speculative positions after the anchor, so a
+training block contains `gamma + 1` positions. `anchors_per_sample` controls how
+many training blocks are sampled from each logical sequence. Blocks never cross
+packed sample boundaries.
+
+## Configure DSpark
+
+DSpark uses the same block-parallel body and adds a Markov head plus an optional
+confidence head. Use `block_size`, not the DFlash-only `gamma` field:
+
+```yaml
+policy:
+  model_name: Qwen/Qwen3-8B
+  megatron_cfg:
+    enabled: true
+    tensor_model_parallel_size: 2
+    pipeline_model_parallel_size: 1
+    context_parallel_size: 1
+    sequence_parallel: false
+    use_fused_linear_logprobs: false
+  dtensor_cfg:
+    enabled: false
+  sequence_packing:
+    enabled: false
+  draft:
+    speculator_type: dspark
+    enabled: true
+    model_name: deepseek-ai/dspark_qwen3_8b_block7
+    model_revision: 03326e5043815da1f81b109078b2889737c26017
+    loss_weight: 1.0
+    block_size: 7
+    anchors_per_sample: 2
+    mask_token_id: 151669
+    target_hidden_state_layer_ids: [1, 9, 17, 25, 33]
+    num_layers: 5
+    markov_rank: 256
+    markov_head_type: vanilla
+    confidence_enabled: true
+    confidence_with_markov: true
+    ce_loss_weight: 0.1
+    tv_loss_weight: 0.9
+    confidence_loss_weight: 1.0
+    loss_decay_gamma: 4.0
+  generation:
+    backend: vllm
+    vllm_kwargs:
+      speculative_config:
+        method: dspark
+        model: ${policy.draft.model_name}
+        revision: ${policy.draft.model_revision}
+        num_speculative_tokens: ${policy.draft.block_size}
+        draft_tensor_parallel_size: 1
+```
+
+DSpark online training uses the live target vocabulary and target-owned
+embedding/output weights. Omit `draft_vocab_size`, or set it to exactly the
+target vocabulary size. A different draft vocabulary is rejected because the
+online remapping path is not supported.
+
+## Enable Packed Context Parallel Training
+
+The following command runs the DFlash recipe with TP2 and CP4 on eight training
+GPUs. It enables sequence packing and target sequence parallelism. Colocated
+vLLM generation uses CP1 and PP1.
+
+```bash
+uv run examples/run_grpo.py \
+  --config examples/configs/recipes/llm/grpo-qwen3-8b-1n8g-megatron-dflash.yaml \
+  data_plane.enabled=true \
+  policy.megatron_cfg.tensor_model_parallel_size=2 \
+  policy.megatron_cfg.pipeline_model_parallel_size=1 \
+  policy.megatron_cfg.context_parallel_size=4 \
+  policy.megatron_cfg.sequence_parallel=true \
+  policy.sequence_packing.enabled=true \
+  policy.make_sequence_length_divisible_by=16
+```
+
+For TP2 x CP2, set `context_parallel_size=2` and
+`make_sequence_length_divisible_by=8`. Set the latter to a multiple of
+`2 * TP * CP`. Keep packing enabled for CP2 or CP4. Target sequence
+parallelism is optional with packing, but requires packing when enabled.
+
+Packed sequences carry sample IDs and sequence boundaries. NeMo RL uses them
+to reconstruct target sequence-parallel captures, gather projected key/value
+tensors across CP, and assign each draft window to one CP rank. Do not add
+layout fields to the input data.
+
+## Loss, Diagnostics, and Optimizer Settings
+
+The total objective is the policy loss plus `policy.draft.loss_weight` times the
+provider-specific draft loss. NeMo RL logs the normalized value as
+`train/draft_loss`.
+
+Set `policy.draft.update_probe_enabled=true` for a short diagnostic run. After
+the optimizer step, NeMo RL checks that the draft model had nonzero gradients
+and that at least one draft parameter changed. Disable the probe for ordinary
+training.
+
+To use a separate learning-rate schedule for draft parameters, set:
+
+```yaml
+policy:
+  draft:
+    optimizer:
+      lr: 1.0e-5
+      min_lr: 1.0e-6
+      weight_decay: 0.01
+```
+
+If `optimizer` is `null`, draft parameters use the policy optimizer settings.
+
+## Checkpoint and Refit Behavior
+
+- `policy.draft.model_name` initializes the trainer-owned draft model.
+- Training checkpoints include the draft model and its optimizer state when
+  draft training is enabled.
+- `policy.generation.vllm_kwargs.speculative_config.model` initializes the vLLM
+  drafter before the first refit.
+- Each successful refit updates both target and draft weights. The generation
+  `method` must match `policy.draft.speculator_type`.
+- DFlash and DSpark generation requires PP1. Keep generation CP at 1 even when
+  policy training uses CP2 or CP4.
+
+## Common Configuration Errors
+
+- `context_parallel_size > 1` with packing disabled
+- `context_parallel_size > 1` without `data_plane.enabled=true`, or with
+  `grpo.async_grpo.enabled=true`
+- async GRPO with `data_plane.enabled=false`
+- target sequence parallelism with packing disabled
+- target TP other than 2, or training PP/virtual PP greater than 1
+- generation CP or PP greater than 1
+- `use_fused_linear_logprobs=true`
+- a generation method that does not match the training speculator
+- `gamma` in a DSpark block or `block_size` in a DFlash block
+- a DSpark `draft_vocab_size` different from the live target vocabulary
+
+NeMo RL rejects these combinations during configuration or worker setup.

docs/guides/refit.md

@@ -67,6 +67,16 @@ workers can send weights. Sparse delta is currently limited to GRPO. NIXL is
 initialized by the GRPO and distillation setup paths; PPO currently requires
 colocated generation.
 
+### Draft Weights with Context Parallel Training
+
+With `refit_transport: nccl_reshard`, a Megatron trainer can send draft weights
+from its TP, PP, and CP lanes to vLLM. The vLLM generation layout must use CP1.
+Setup rejects a `decode_context_parallel_size` or `prefill_context_parallel_size`
+value other than 1 in `vllm_kwargs` or `vllm_cfg`.
+This transport support does not, by itself, enable a draft-training layout;
+configure a supported training path separately. The vLLM reload API does not
+support trainer-owned draft weights.
+
 ## Minimal Configuration
 
 Colocated vLLM and SGLang refit need no transport configuration:

docs/index.md

@@ -128,6 +128,13 @@ Create custom reward environments and integrate them with NeMo RL training pipel
 Configure offline and online Eagle3 draft-model workflows to accelerate rollout generation with vLLM.
 :::
 
+:::{grid-item-card} {octicon}`iterations` DFlash and DSpark Co-Training
+:link: guides/dflash-dspark-speculative-decoding
+:link-type: doc
+
+Train block-parallel DFlash and DSpark draft models with packed TP2 and context-parallel Megatron policies.
+:::
+
 :::{grid-item-card} {octicon}`unmute` Audio GRPO on AVQA
 :link: guides/grpo-audio
 :link-type: doc
@@ -334,6 +341,7 @@ guides/async-grpo.md
 guides/single-controller.md
 guides/quantization-aware-rl.md
 guides/eagle3-speculative-decoding.md
+guides/dflash-dspark-speculative-decoding.md
 guides/yarn-long-context.md
 guides/xtoken-off-policy-distillation.md
 guides/refit.md
@@ -398,6 +406,8 @@ design-docs/dspark-auxiliary-heads.md
 design-docs/media-token-validity-mask.md
 design-docs/automodel-context-parallel.md
 design-docs/token-capture-ledger.md
+design-docs/dflash-draft-body.md
+design-docs/dspark-objectives.md
 ```
 
 ```{toctree}

examples/configs/grpo_math_1B.yaml

@@ -363,13 +363,17 @@ policy:
     env_vars: null
 
   draft:
-    # Draft co-training speculator. Currently only eagle3 is supported.
+    # Default draft speculator; DFlash and DSpark have dedicated recipe configs.
     speculator_type: eagle3
     enabled: false
     model_name: null
     loss_weight: 0.1
     num_layers: null
     aux_layer_indices: null
+    # Optional draft-only optimizer schedule. Set a mapping with `lr` and
+    # optional `min_lr`/`weight_decay` to enable it.
+    optimizer: null
+
     # Tokens per FP32 vocab tile in the streaming draft soft-CE. Lower it if the
     # tile (token_chunk_size x vocab, in fp32) is too large for the device;
     # raise it if the GPU is launch-bound.