From 589295c7d8b0b4a1e7a23565b2e0c9ed0f14c911 Mon Sep 17 00:00:00 2001
From: SlyneD <slyned@nvidia.com>
Date: Tue, 18 Aug 2026 13:49:34 -0700
Subject: [PATCH 1/8] feat(speechlm2): add DFlash training for SALM

Signed-off-by: SlyneD <slyned@nvidia.com>
---
 examples/speechlm2/conf/salm_automodel.yaml   |  19 +
 examples/speechlm2/salm_train.py              |  18 +
 nemo/collections/speechlm2/parts/dflash.py    | 508 ++++++++++++++++++
 .../collections/speechlm2/test_salm_dflash.py | 229 ++++++++
 4 files changed, 774 insertions(+)
 create mode 100644 nemo/collections/speechlm2/parts/dflash.py
 create mode 100644 tests/collections/speechlm2/test_salm_dflash.py

diff --git a/examples/speechlm2/conf/salm_automodel.yaml b/examples/speechlm2/conf/salm_automodel.yaml
index 57377510f81c..3cd6722a1f0c 100644
--- a/examples/speechlm2/conf/salm_automodel.yaml
+++ b/examples/speechlm2/conf/salm_automodel.yaml
@@ -290,3 +290,22 @@ exp_manager:
      save_top_k: 1
      always_save_nemo: false
      save_nemo_on_train_end: false
+
+# Optional DFlash draft training from a frozen audio-conditioned SALMAutomodel.
+# This initial integration supports BSHD batches with pp_size=cp_size=1. Keep
+# num_anchors small: Automodel's generic DFlash trainer currently materializes
+# draft vocabulary logits.
+dflash:
+  enabled: false
+  mask_token_id: null  # Set this to an unused token in the target vocabulary.
+  draft_num_hidden_layers: 2
+  target_layer_ids: null
+  draft_model_config: {}
+  block_size: 8
+  num_anchors: 1
+  loss_decay_gamma: 7.0
+  attention_backend: sdpa
+  activation_checkpointing: false
+  target_dtype: bfloat16
+  lr: 6.0e-4
+  output_dir: ./outputs/salm_dflash
diff --git a/examples/speechlm2/salm_train.py b/examples/speechlm2/salm_train.py
index 04e47d1b0c81..1e1607b96567 100644
--- a/examples/speechlm2/salm_train.py
+++ b/examples/speechlm2/salm_train.py
@@ -43,6 +43,8 @@ def train(cfg):
         torch.distributed.init_process_group(backend="nccl")
     seed_everything(cfg.data.train_ds.seed)
     torch.set_float32_matmul_precision("medium")
+    if cfg.get("dflash", {}).get("enabled", False) and not cfg.model.get("use_nemo_automodel", False):
+        raise ValueError("SALM DFlash training requires model.use_nemo_automodel=true")
     trainer = Trainer(**resolve_trainer_cfg(cfg.trainer))
     log_dir = exp_manager(trainer, cfg.get("exp_manager", None))
     # Insert at position 0 so our ``on_train_batch_end`` runs BEFORE the
@@ -52,6 +54,22 @@ def train(cfg):
     trainer.callbacks.insert(0, TrainingStatsCallback())
     OmegaConf.save(cfg, log_dir / "exp_config.yaml")
 
+    if cfg.get("dflash", {}).get("enabled", False):
+        from nemo.collections.speechlm2 import SALMAutomodel
+        from nemo.collections.speechlm2.parts.dflash import SALMDFlashModule
+
+        model_cfg = OmegaConf.to_container(cfg.model, resolve=True)
+        model_cfg["torch_dtype"] = cfg.dflash.get("target_dtype", "bfloat16")
+        with trainer.init_module():
+            target_model = SALMAutomodel(model_cfg)
+        model = SALMDFlashModule(target_model, OmegaConf.to_container(cfg, resolve=True))
+        dataset = _create_salm_dataset(target_model.tokenizer, cfg.data)
+        datamodule = DataModule(cfg.data, tokenizer=target_model.tokenizer, dataset=dataset)
+        trainer.fit(model, datamodule)
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()
+        return
+
     model_cls = SALM
     if cfg.model.get("use_nemo_automodel", False):
         from nemo.collections.speechlm2 import SALMAutomodel
diff --git a/nemo/collections/speechlm2/parts/dflash.py b/nemo/collections/speechlm2/parts/dflash.py
new file mode 100644
index 000000000000..2bda5ce46c88
--- /dev/null
+++ b/nemo/collections/speechlm2/parts/dflash.py
@@ -0,0 +1,508 @@
+# 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.
+
+"""DFlash draft training for audio-conditioned SALM targets."""
+
+from __future__ import annotations
+
+import inspect
+from collections import defaultdict
+from collections.abc import Sequence
+
+import torch
+from lightning import LightningModule
+from nemo_automodel.components.speculative.dflash.core import DFlashTrainerModule, NoValidAnchorsError
+from nemo_automodel.components.speculative.dflash.draft_qwen3 import Qwen3DFlashDraftModel, build_target_layer_ids
+from torch import nn
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+
+from nemo.collections.speechlm2.models.salm import replace_placeholders_and_build_targets
+from nemo.collections.speechlm2.parts.cp_helpers import encode_audio_with_cp_distribution, get_perception_fsdp_group
+from nemo.core.classes.common import safe_instantiate
+
+_DRAFT_CONFIG_MANAGED_KEYS = {
+    "architectures",
+    "block_size",
+    "dflash_config",
+    "layer_types",
+    "max_window_layers",
+    "num_hidden_layers",
+    "num_target_layers",
+}
+
+
+def _all_ranks_agree(local_condition: bool, device: torch.device) -> bool:
+    """Return whether every distributed rank reports a true condition."""
+    if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+        return local_condition
+    flag = torch.tensor([int(local_condition)], device=device, dtype=torch.int32)
+    torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MIN)
+    return bool(flag.item())
+
+
+def _max_rank_value(local_value: int, device: torch.device) -> int:
+    """Return the maximum integer reported by any distributed rank."""
+    if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+        return local_value
+    value = torch.tensor([local_value], device=device, dtype=torch.int32)
+    torch.distributed.all_reduce(value, op=torch.distributed.ReduceOp.MAX)
+    return int(value.item())
+
+
+def _all_ranks_report_same_value(local_value: int, device: torch.device) -> bool:
+    """Return whether every distributed rank reports the same integer."""
+    if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+        return True
+    extrema = torch.tensor([local_value, -local_value], device=device, dtype=torch.int32)
+    torch.distributed.all_reduce(extrema, op=torch.distributed.ReduceOp.MIN)
+    return int(extrema[0].item()) == -int(extrema[1].item())
+
+
+def _preprocessing_signature(batch: dict[str, torch.Tensor]) -> int:
+    """Describe rank-local branches that may enter distributed perception."""
+    audio_lens = batch.get("audio_lens")
+    has_audio = audio_lens is not None and audio_lens.numel() > 0
+    has_speaker_targets = batch.get("spk_targets") is not None
+    return int(has_audio) | (int(has_speaker_targets) << 1)
+
+
+def _synchronize_ep_group_before_target_forward(moe_mesh) -> None:
+    """Keep rank-local audio preprocessing skew out of DeepEP's timeout."""
+    if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+        return
+    if moe_mesh is None or "ep" not in moe_mesh.mesh_dim_names:
+        return
+    ep_mesh = moe_mesh["ep"]
+    if ep_mesh.size() > 1:
+        torch.distributed.barrier(group=ep_mesh.get_group())
+
+
+def _has_valid_dflash_anchors(loss_mask: torch.Tensor, block_size: int) -> bool:
+    """Return whether this rank can form at least one DFlash anchor."""
+    max_anchor = max(loss_mask.shape[1] - block_size, 0)
+    return bool((loss_mask[:, : max_anchor + 1] > 0.5).any().item())
+
+
+def _build_draft_config(
+    target_config,
+    dflash_config: dict,
+    block_size: int,
+    mask_token_id: int,
+) -> tuple[Qwen3Config, list[int]]:
+    """Create the Qwen3-shaped DFlash draft config for a SALM target."""
+    num_target_layers = int(target_config.num_hidden_layers)
+    draft_layers = int(dflash_config.get("draft_num_hidden_layers", 2))
+    target_layer_ids = list(
+        dflash_config.get("target_layer_ids") or build_target_layer_ids(num_target_layers, draft_layers)
+    )
+    if len(set(target_layer_ids)) != len(target_layer_ids):
+        raise ValueError("dflash.target_layer_ids must be unique")
+    if not target_layer_ids or min(target_layer_ids) < 0 or max(target_layer_ids) >= num_target_layers:
+        raise ValueError(f"dflash.target_layer_ids must be within [0, {num_target_layers})")
+
+    architecture = dict(dflash_config.get("draft_model_config") or {})
+    managed = sorted(_DRAFT_CONFIG_MANAGED_KEYS.intersection(architecture))
+    if managed:
+        raise ValueError(f"dflash.draft_model_config cannot override managed keys: {', '.join(managed)}")
+
+    draft_dict = target_config.to_dict()
+    draft_dict.update(architecture)
+    draft_dict.update(
+        {
+            "architectures": ["Qwen3DFlashDraftModel"],
+            "num_hidden_layers": draft_layers,
+            "layer_types": ["full_attention"] * draft_layers,
+            "max_window_layers": draft_layers,
+            "num_target_layers": num_target_layers,
+            "block_size": block_size,
+            "dflash_config": {
+                "mask_token_id": mask_token_id,
+                "target_layer_ids": target_layer_ids,
+            },
+        }
+    )
+    draft_config = Qwen3Config.from_dict(draft_dict)
+    if draft_config.hidden_size != target_config.hidden_size:
+        raise ValueError(
+            "The DFlash draft hidden_size must match the frozen target because its embeddings and LM head are reused "
+            f"({draft_config.hidden_size} != {target_config.hidden_size})."
+        )
+    return draft_config, target_layer_ids
+
+
+def _expand_ids_with_audio(
+    input_ids: torch.Tensor,
+    replacements: Sequence[torch.Tensor],
+    padding_id: int,
+    placeholder_id: int,
+    mask_token_id: int,
+) -> torch.Tensor:
+    """Expand audio placeholders to mask-token runs matching fused embeddings."""
+    rows = []
+    replacement_idx = 0
+    for row in input_ids:
+        non_padding = (row != padding_id).nonzero(as_tuple=False)
+        first_non_padding = int(non_padding[0]) if non_padding.numel() else row.numel()
+        row = row[first_non_padding:]
+        pieces = []
+        for token in row:
+            if int(token) == placeholder_id:
+                length = replacements[replacement_idx].shape[0]
+                replacement_idx += 1
+                pieces.append(torch.full((length,), mask_token_id, dtype=row.dtype, device=row.device))
+            else:
+                pieces.append(token.view(1))
+        rows.append(torch.cat(pieces) if pieces else row)
+    if replacement_idx != len(replacements):
+        raise ValueError(f"Used {replacement_idx} of {len(replacements)} audio replacements")
+
+    max_len = max(row.numel() for row in rows)
+    expanded = torch.full((len(rows), max_len), padding_id, dtype=input_ids.dtype, device=input_ids.device)
+    for index, row in enumerate(rows):
+        expanded[index, -row.numel() :] = row
+    return expanded
+
+
+def _get_consolidated_model_state_dict(model: nn.Module) -> dict[str, torch.Tensor]:
+    """Gather an FSDP2 draft into a Hugging Face-saveable rank-zero state dict."""
+    if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+        return model.state_dict()
+
+    from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
+
+    return get_model_state_dict(model, options=StateDictOptions(full_state_dict=True, cpu_offload=True))
+
+
+class SALMDFlashModule(LightningModule):
+    """Train a Qwen3-style DFlash draft from a frozen ``SALMAutomodel`` target."""
+
+    _CHECKPOINT_STATE_PREFIX = "draft_model."
+    _REBUILT_STATE_PREFIXES = ("target.", "trainer_module.")
+
+    def __init__(self, target_model: nn.Module, cfg: dict):
+        super().__init__()
+        self.target = target_model
+        self.cfg = cfg
+        self.dflash_config = cfg.get("dflash", cfg)
+        self.block_size = int(self.dflash_config.get("block_size", 8))
+        mask_token_id = self.dflash_config.get("mask_token_id")
+        if mask_token_id is None:
+            raise ValueError("dflash.mask_token_id must identify an unused token in the target vocabulary")
+        self.mask_token_id = int(mask_token_id)
+        self.attention_backend = str(self.dflash_config.get("attention_backend", "sdpa"))
+        self.output_dir = self.dflash_config.get("output_dir")
+        self.learning_rate = float(self.dflash_config.get("lr", 6e-4))
+        self.draft_model = None
+        self.trainer_module = None
+        self.target_layer_ids = None
+        self._partial_val_metrics = defaultdict(list)
+        self.register_state_dict_post_hook(self._keep_draft_checkpoint_state)
+
+    @staticmethod
+    def _keep_draft_checkpoint_state(module, state_dict, prefix, local_metadata) -> None:
+        """Exclude the rebuilt frozen target from resumable checkpoints."""
+        draft_prefix = f"{prefix}{module._CHECKPOINT_STATE_PREFIX}"
+        for key in tuple(state_dict):
+            if not key.startswith(draft_prefix):
+                del state_dict[key]
+
+    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
+        """Load a draft-only checkpoint while allowing the target to be rebuilt."""
+        incompatible = super().load_state_dict(state_dict, strict=False, assign=assign)
+        missing = [key for key in incompatible.missing_keys if not key.startswith(self._REBUILT_STATE_PREFIXES)]
+        if strict and (missing or incompatible.unexpected_keys):
+            raise RuntimeError(
+                f"Error loading {type(self).__name__}: missing={missing}, unexpected={incompatible.unexpected_keys}"
+            )
+        return type(incompatible)(missing, incompatible.unexpected_keys)
+
+    def configure_model(self) -> None:
+        """Build and shard the frozen SALM target before creating the draft."""
+        if self.draft_model is not None:
+            return
+
+        strategy = self.trainer.strategy
+        distributed_setup = getattr(strategy, "distributed_setup", None)
+        if distributed_setup is not None:
+            mesh_context = distributed_setup.mesh_context
+            if mesh_context.pp_size > 1 or mesh_context.cp_size > 1:
+                raise NotImplementedError("SALM DFlash currently requires pp_size=cp_size=1")
+        self.target._trainer = self.trainer
+        self.target.configure_model(
+            distributed_setup=distributed_setup,
+            activation_checkpointing_perception=getattr(strategy, "activation_checkpointing_perception", False),
+        )
+        self.target.eval()
+        self.target.requires_grad_(False)
+
+        target_config = self.target.llm.config
+        vocab_size = int(target_config.vocab_size)
+        if not 0 <= self.mask_token_id < vocab_size:
+            raise ValueError(f"dflash.mask_token_id={self.mask_token_id} is outside [0, {vocab_size})")
+        draft_config, self.target_layer_ids = _build_draft_config(
+            target_config, self.dflash_config, self.block_size, self.mask_token_id
+        )
+        draft_config._attn_implementation = self.attention_backend
+        dtype = next(self.target.llm.parameters()).dtype
+        self.draft_model = Qwen3DFlashDraftModel(draft_config).to(self.target.device, dtype=dtype)
+        if self.dflash_config.get("activation_checkpointing", False):
+            self.draft_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+
+        self.trainer_module = DFlashTrainerModule(
+            draft_model=self.draft_model,
+            target_lm_head=self.target.llm.get_output_embeddings(),
+            target_embed_tokens=self.target.llm.get_input_embeddings(),
+            mask_token_id=self.mask_token_id,
+            block_size=self.block_size,
+            attention_backend=self.attention_backend,
+            num_anchors=int(self.dflash_config.get("num_anchors", 1)),
+            loss_decay_gamma=self.dflash_config.get("loss_decay_gamma", 7.0),
+            loss_type=str(self.dflash_config.get("loss_type", "dflash")),
+            prefix_weight_base=float(self.dflash_config.get("prefix_weight_base", 0.9)),
+        )
+
+        device_mesh = self.device_mesh
+        dim_names = device_mesh.mesh_dim_names
+        if "dp_replicate" in dim_names and "dp_shard_cp" in dim_names:
+            draft_fsdp_mesh = device_mesh["dp_replicate", "dp_shard_cp"]
+        elif "dp_shard_cp" in dim_names:
+            draft_fsdp_mesh = device_mesh["dp_shard_cp"]
+        else:
+            draft_fsdp_mesh = device_mesh["dp"]
+        if draft_fsdp_mesh.size() > 1:
+            from torch.distributed.fsdp import fully_shard
+
+            self.draft_model = fully_shard(self.draft_model, mesh=draft_fsdp_mesh)
+
+        if any(parameter.requires_grad for parameter in self.target.parameters()):
+            raise RuntimeError("The DFlash SALM target must be fully frozen")
+
+    @property
+    def device(self):
+        return next(self.draft_model.parameters()).device
+
+    def _audio_embeddings(self, batch: dict[str, torch.Tensor]) -> list[torch.Tensor]:
+        spk_targets = batch.get("spk_targets")
+        if self.target._uses_parallel_expert_encoder() and spk_targets is None:
+            embeddings, lengths = self.target.perception(
+                input_signal=batch["audios"], input_signal_length=batch["audio_lens"]
+            )
+            return [embedding[:length] for embedding, length in zip(embeddings, lengths)]
+
+        device_mesh = getattr(self.target, "_device_mesh", None)
+        return encode_audio_with_cp_distribution(
+            self.target.perception,
+            batch["audios"],
+            batch["audio_lens"],
+            chunk_size_seconds=self.target.cfg.get("encoder_chunk_size_seconds"),
+            sampling_rate=self.target.sampling_rate,
+            cp_mesh=None,
+            spk_targets=spk_targets,
+            fsdp_sync_group=get_perception_fsdp_group(device_mesh),
+        )
+
+    def _prepare_batch(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        if self.target.cfg.get("packed_sequences", False):
+            raise NotImplementedError("SALM DFlash currently requires model.packed_sequences=false")
+
+        input_ids = batch["input_ids"]
+        audio_embeddings = self._audio_embeddings(batch)
+        text_ids = torch.where(input_ids == self.target.audio_locator_tag_id, 0, input_ids)
+        text_embeddings = self.target._embed_tokens(text_ids)
+        target_ids = input_ids.where(batch["loss_mask"], -100)
+        input_embeddings, target_ids, attention_mask = replace_placeholders_and_build_targets(
+            input_ids=input_ids,
+            embeds=text_embeddings,
+            padding_id=self.target.text_pad_id,
+            placeholder_id=self.target.audio_locator_tag_id,
+            replacements=audio_embeddings,
+            target_ids=target_ids,
+        )
+        expanded_ids = _expand_ids_with_audio(
+            input_ids,
+            audio_embeddings,
+            self.target.text_pad_id,
+            self.target.audio_locator_tag_id,
+            self.mask_token_id,
+        )
+        return {
+            "input_ids": expanded_ids[:, :-1],
+            "input_embeddings": input_embeddings[:, :-1],
+            "attention_mask": attention_mask[:, :-1],
+            "loss_mask": target_ids[:, 1:].ne(-100),
+        }
+
+    @torch.no_grad()
+    def _target_hidden_states(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor:
+        """Run the frozen audio-conditioned target and concatenate configured layers."""
+        if hasattr(self.target.llm, "model") and hasattr(self.target.llm.model, "layers"):
+            layer_container = self.target.llm.model.layers
+        elif hasattr(self.target.llm, "layers"):
+            layer_container = self.target.llm.layers
+        elif hasattr(self.target.llm, "transformer") and hasattr(self.target.llm.transformer, "h"):
+            layer_container = self.target.llm.transformer.h
+        else:
+            raise ValueError("Unsupported SALM target structure for DFlash hidden-state capture")
+        if isinstance(layer_container, nn.ModuleDict):
+            layers = [layer_container[str(index)] for index in range(len(layer_container))]
+        else:
+            layers = list(layer_container)
+
+        captured = {}
+        handles = []
+
+        def make_hook(layer_id: int):
+            def hook(_module, _args, output):
+                captured[layer_id] = output[0] if isinstance(output, tuple) else output
+
+            return hook
+
+        for layer_id in self.target_layer_ids:
+            handles.append(layers[layer_id].register_forward_hook(make_hook(layer_id)))
+
+        forward_kwargs = {
+            "inputs_embeds": inputs["input_embeddings"],
+            "attention_mask": inputs["attention_mask"],
+            "output_hidden_states": False,
+            "use_cache": False,
+            "return_dict": True,
+        }
+        if "compute_logits" in inspect.signature(type(self.target.llm).forward).parameters:
+            forward_kwargs["compute_logits"] = False
+        try:
+            self.target.llm(**forward_kwargs)
+        finally:
+            for handle in handles:
+                handle.remove()
+        if len(captured) != len(self.target_layer_ids):
+            raise RuntimeError(f"Expected {len(self.target_layer_ids)} captured target layers, got {sorted(captured)}")
+        return torch.cat([captured[layer_id] for layer_id in self.target_layer_ids], dim=-1)
+
+    def _run_batch(self, batch: dict[str, torch.Tensor]):
+        inputs = self._prepare_batch(batch)
+        if not _all_ranks_agree(
+            _has_valid_dflash_anchors(inputs["loss_mask"], self.block_size), inputs["loss_mask"].device
+        ):
+            raise NoValidAnchorsError("At least one rank has no valid DFlash anchors")
+        _synchronize_ep_group_before_target_forward(getattr(self.trainer.strategy, "moe_mesh", None))
+        hidden_states = self._target_hidden_states(inputs)
+        return self.trainer_module(
+            input_ids=inputs["input_ids"],
+            hidden_states=hidden_states,
+            loss_mask=inputs["loss_mask"],
+        )
+
+    def training_step(self, batch, batch_idx):
+        batches = list(batch.values()) if isinstance(batch, dict) and "input_ids" not in batch else [batch]
+        losses = []
+        num_batches = _max_rank_value(len(batches), self.device)
+        for dataset_index in range(num_batches):
+            dataset_batch = batches[dataset_index] if dataset_index < len(batches) else None
+            if not _all_ranks_agree(dataset_batch is not None, self.device):
+                continue
+            assert dataset_batch is not None
+            signature = _preprocessing_signature(dataset_batch)
+            if not _all_ranks_report_same_value(signature, self.device):
+                continue
+            try:
+                metrics = self._run_batch(dataset_batch)
+            except NoValidAnchorsError:
+                continue
+            losses.append(metrics.loss)
+            self.log("train/dflash_loss", metrics.loss, on_step=True, prog_bar=True)
+            self.log("train/dflash_accuracy", metrics.accuracy, on_step=True)
+            self.log("train/accept_len", metrics.accept_len, on_step=True)
+        if not losses:
+            return torch.zeros((), device=self.device, requires_grad=True)
+        return torch.stack(losses).mean()
+
+    def on_validation_epoch_start(self) -> None:
+        self._partial_val_metrics.clear()
+
+    def validation_step(self, batch, batch_idx) -> None:
+        batches = (
+            list(batch.items()) if isinstance(batch, dict) and "input_ids" not in batch else [("validation", batch)]
+        )
+        num_batches = _max_rank_value(len(batches), self.device)
+        for dataset_index in range(num_batches):
+            dataset_name, dataset_batch = batches[dataset_index] if dataset_index < len(batches) else ("missing", None)
+            if not _all_ranks_agree(dataset_batch is not None, self.device):
+                continue
+            assert dataset_batch is not None
+            signature = _preprocessing_signature(dataset_batch)
+            if not _all_ranks_report_same_value(signature, self.device):
+                continue
+            try:
+                metrics = self._run_batch(dataset_batch)
+            except NoValidAnchorsError:
+                continue
+            metric_dtype = metrics.loss.dtype
+            metric_device = metrics.loss.device
+            self._partial_val_metrics[dataset_name].append(
+                torch.stack(
+                    [
+                        metrics.loss.detach() * metrics.loss_weight.to(metric_dtype),
+                        metrics.loss_weight.to(dtype=metric_dtype, device=metric_device),
+                        metrics.correct_tokens.to(dtype=metric_dtype, device=metric_device),
+                        metrics.valid_tokens.to(dtype=metric_dtype, device=metric_device),
+                        metrics.accept_len_sum.to(dtype=metric_dtype, device=metric_device),
+                        metrics.valid_blocks.to(dtype=metric_dtype, device=metric_device),
+                    ]
+                )
+            )
+
+    def on_validation_epoch_end(self) -> None:
+        all_sums = []
+        for dataset_name, partial_metrics in self._partial_val_metrics.items():
+            if not partial_metrics:
+                continue
+            metric_sums = torch.stack(partial_metrics).sum(dim=0)
+            if torch.distributed.is_available() and torch.distributed.is_initialized():
+                torch.distributed.all_reduce(metric_sums, op=torch.distributed.ReduceOp.SUM)
+            all_sums.append(metric_sums)
+            self._log_validation_metrics(metric_sums, suffix=f"/{dataset_name}")
+        if all_sums:
+            self._log_validation_metrics(torch.stack(all_sums).sum(dim=0))
+        self._partial_val_metrics.clear()
+
+    def _log_validation_metrics(self, metric_sums: torch.Tensor, suffix: str = "") -> None:
+        loss_sum, loss_weight, correct, valid, accept_sum, valid_blocks = metric_sums
+        self.log(f"val/dflash_loss{suffix}", loss_sum / loss_weight.clamp_min(1), on_epoch=True)
+        self.log(f"val/dflash_accuracy{suffix}", correct / valid.clamp_min(1), on_epoch=True)
+        self.log(f"val/accept_len{suffix}", accept_sum / valid_blocks.clamp_min(1), on_epoch=True)
+
+    def configure_optimizers(self):
+        optimizer_config = self.dflash_config.get("optimizer")
+        if optimizer_config is None:
+            return torch.optim.AdamW(self.draft_model.parameters(), lr=self.learning_rate)
+        optimizer = safe_instantiate(
+            optimizer_config,
+            params=self.draft_model.parameters(),
+            _convert_="all",
+        )
+        scheduler_config = self.dflash_config.get("lr_scheduler")
+        if scheduler_config is None:
+            return optimizer
+        scheduler = safe_instantiate(scheduler_config, optimizer=optimizer, _convert_="all")
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {"scheduler": scheduler, "interval": "step", "frequency": 1},
+        }
+
+    def on_train_end(self) -> None:
+        if not self.output_dir:
+            return
+        state_dict = _get_consolidated_model_state_dict(self.draft_model)
+        if self.trainer.is_global_zero:
+            self.draft_model.save_pretrained(self.output_dir, state_dict=state_dict)
diff --git a/tests/collections/speechlm2/test_salm_dflash.py b/tests/collections/speechlm2/test_salm_dflash.py
new file mode 100644
index 000000000000..83a42432d168
--- /dev/null
+++ b/tests/collections/speechlm2/test_salm_dflash.py
@@ -0,0 +1,229 @@
+# 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 types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from torch import nn
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+
+from nemo.collections.speechlm2.parts import dflash as salm_dflash
+
+
+class _FakeMoEMesh:
+    mesh_dim_names = ("ep_shard", "ep")
+
+    def __init__(self, ep_mesh):
+        self.ep_mesh = ep_mesh
+
+    def __getitem__(self, name):
+        assert name == "ep"
+        return self.ep_mesh
+
+
+def test_synchronize_ep_group_uses_ep_process_group(monkeypatch):
+    group = object()
+    ep_mesh = SimpleNamespace(size=lambda: 8, get_group=lambda: group)
+    calls = []
+    monkeypatch.setattr(salm_dflash.torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(salm_dflash.torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(salm_dflash.torch.distributed, "barrier", lambda *, group: calls.append(group))
+
+    salm_dflash._synchronize_ep_group_before_target_forward(_FakeMoEMesh(ep_mesh))
+
+    assert calls == [group]
+
+
+def test_synchronize_ep_group_is_noop_without_distributed_ep(monkeypatch):
+    calls = []
+    monkeypatch.setattr(salm_dflash.torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(salm_dflash.torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(salm_dflash.torch.distributed, "barrier", lambda **kwargs: calls.append(kwargs))
+
+    salm_dflash._synchronize_ep_group_before_target_forward(None)
+    salm_dflash._synchronize_ep_group_before_target_forward(
+        _FakeMoEMesh(SimpleNamespace(size=lambda: 1, get_group=lambda: object()))
+    )
+
+    assert calls == []
+
+
+def test_expand_ids_with_audio_preserves_internal_pad_valued_tokens():
+    input_ids = torch.tensor([[0, 0, 11, 99, 0, 12]])
+    audio_embeddings = [torch.randn(3, 4)]
+
+    expanded = salm_dflash._expand_ids_with_audio(
+        input_ids,
+        audio_embeddings,
+        padding_id=0,
+        placeholder_id=99,
+        mask_token_id=18,
+    )
+
+    assert expanded.tolist() == [[11, 18, 18, 18, 0, 12]]
+
+
+def test_expand_ids_with_audio_left_pads_rows_to_common_length():
+    input_ids = torch.tensor([[0, 10, 99, 12], [20, 21, 22, 23]])
+    audio_embeddings = [torch.randn(2, 4)]
+
+    expanded = salm_dflash._expand_ids_with_audio(
+        input_ids,
+        audio_embeddings,
+        padding_id=0,
+        placeholder_id=99,
+        mask_token_id=18,
+    )
+
+    assert expanded.tolist() == [[10, 18, 18, 12], [20, 21, 22, 23]]
+
+
+def test_expand_ids_with_audio_requires_every_replacement_to_be_used():
+    with pytest.raises(ValueError, match="Used 0 of 1"):
+        salm_dflash._expand_ids_with_audio(
+            torch.tensor([[1, 2, 3]]),
+            [torch.randn(2, 4)],
+            padding_id=0,
+            placeholder_id=99,
+            mask_token_id=18,
+        )
+
+
+def test_build_draft_config_applies_explicit_architecture_and_layer_taps():
+    target_config = Qwen3Config(
+        hidden_size=64,
+        intermediate_size=128,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        num_hidden_layers=8,
+        head_dim=16,
+        vocab_size=128,
+    )
+
+    draft_config, layer_ids = salm_dflash._build_draft_config(
+        target_config,
+        {
+            "draft_num_hidden_layers": 2,
+            "target_layer_ids": [1, 6],
+            "draft_model_config": {"intermediate_size": 96},
+        },
+        block_size=8,
+        mask_token_id=18,
+    )
+
+    assert layer_ids == [1, 6]
+    assert draft_config.num_hidden_layers == 2
+    assert draft_config.intermediate_size == 96
+    assert draft_config.dflash_config == {"mask_token_id": 18, "target_layer_ids": [1, 6]}
+
+
+def test_build_draft_config_rejects_managed_overrides():
+    target_config = Qwen3Config(hidden_size=64, num_attention_heads=4, num_hidden_layers=8, vocab_size=128)
+
+    with pytest.raises(ValueError, match="cannot override managed keys: block_size"):
+        salm_dflash._build_draft_config(
+            target_config,
+            {"draft_model_config": {"block_size": 32}},
+            block_size=8,
+            mask_token_id=18,
+        )
+
+
+class _TargetLLM(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(1))
+        self.layers = nn.ModuleList([nn.Identity(), nn.Identity(), nn.Identity()])
+        self.calls = []
+
+    def forward(
+        self,
+        *,
+        inputs_embeds,
+        attention_mask,
+        output_hidden_states,
+        use_cache,
+        return_dict,
+        compute_logits=True,
+    ):
+        self.calls.append(
+            {
+                "attention_mask": attention_mask,
+                "output_hidden_states": output_hidden_states,
+                "use_cache": use_cache,
+                "return_dict": return_dict,
+                "compute_logits": compute_logits,
+            }
+        )
+        hidden = inputs_embeds
+        for index, layer in enumerate(self.layers, start=1):
+            hidden = layer(hidden + index)
+        return SimpleNamespace(hidden_states=(hidden,))
+
+
+class _TargetModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.llm = _TargetLLM()
+
+
+def test_target_hidden_states_uses_audio_embeddings_and_skips_logits():
+    module = salm_dflash.SALMDFlashModule(_TargetModel(), {"dflash": {"mask_token_id": 18}})
+    module.target_layer_ids = [0, 2]
+    inputs = {
+        "input_embeddings": torch.randn(2, 5, 4),
+        "attention_mask": torch.ones(2, 5, dtype=torch.bool),
+    }
+
+    hidden = module._target_hidden_states(inputs)
+
+    assert hidden.shape == (2, 5, 8)
+    assert torch.allclose(hidden[..., :4], inputs["input_embeddings"] + 1)
+    assert torch.allclose(hidden[..., 4:], inputs["input_embeddings"] + 6)
+    assert module.target.llm.calls == [
+        {
+            "attention_mask": inputs["attention_mask"],
+            "output_hidden_states": False,
+            "use_cache": False,
+            "return_dict": True,
+            "compute_logits": False,
+        }
+    ]
+
+
+def test_get_consolidated_state_dict_uses_plain_state_dict_without_distributed(monkeypatch):
+    expected = {"weight": torch.tensor([1.0])}
+    model = SimpleNamespace(state_dict=Mock(return_value=expected))
+    monkeypatch.setattr(salm_dflash.torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(salm_dflash.torch.distributed, "is_initialized", lambda: False)
+
+    result = salm_dflash._get_consolidated_model_state_dict(model)
+
+    assert result is expected
+    model.state_dict.assert_called_once_with()
+
+
+def test_state_dict_hook_keeps_only_draft_parameters():
+    module = SimpleNamespace(_CHECKPOINT_STATE_PREFIX="draft_model.")
+    state_dict = {
+        "wrapper.draft_model.layer.weight": torch.ones(1),
+        "wrapper.target.layer.weight": torch.ones(1),
+        "wrapper.trainer_module.loss.weight": torch.ones(1),
+    }
+
+    salm_dflash.SALMDFlashModule._keep_draft_checkpoint_state(module, state_dict, "wrapper.", {})
+
+    assert list(state_dict) == ["wrapper.draft_model.layer.weight"]

From b93533c757afcc809cdff08e85989a81a0cb8605 Mon Sep 17 00:00:00 2001
From: SlyneD <slyned@nvidia.com>
Date: Tue, 18 Aug 2026 14:51:45 -0700
Subject: [PATCH 2/8] fix(speechlm2): match Nemotron Lightning DFlash defaults

Signed-off-by: SlyneD <slyned@nvidia.com>
---
 examples/speechlm2/conf/salm_automodel.yaml   | 30 ++++++++++--
 .../collections/speechlm2/test_salm_dflash.py | 48 +++++++++++++++++++
 2 files changed, 74 insertions(+), 4 deletions(-)

diff --git a/examples/speechlm2/conf/salm_automodel.yaml b/examples/speechlm2/conf/salm_automodel.yaml
index 3cd6722a1f0c..1c2a9713d209 100644
--- a/examples/speechlm2/conf/salm_automodel.yaml
+++ b/examples/speechlm2/conf/salm_automodel.yaml
@@ -292,15 +292,37 @@ exp_manager:
      save_nemo_on_train_end: false
 
 # Optional DFlash draft training from a frozen audio-conditioned SALMAutomodel.
+# The defaults mirror the trainable dense core of
+# nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DFlash: six non-causal
+# Qwen3/GQA layers, 32 Q heads, 2 KV heads, 6144-wide MLPs, and the same six
+# target-layer taps. The public checkpoint is a Model Optimizer W4A16_NVFP4
+# inference artifact and cannot directly initialize these trainable BF16
+# nn.Linear weights; these settings reproduce its architecture, not its weights.
+#
 # This initial integration supports BSHD batches with pp_size=cp_size=1. Keep
 # num_anchors small: Automodel's generic DFlash trainer currently materializes
 # draft vocabulary logits.
 dflash:
   enabled: false
-  mask_token_id: null  # Set this to an unused token in the target vocabulary.
-  draft_num_hidden_layers: 2
-  target_layer_ids: null
-  draft_model_config: {}
+  mask_token_id: 990
+  draft_num_hidden_layers: 6
+  target_layer_ids: [1, 5, 19, 29, 41, 51]
+  draft_model_config:
+    attention_bias: false
+    attention_dropout: 0.0
+    head_dim: 128
+    hidden_act: silu
+    intermediate_size: 6144
+    max_position_embeddings: 1048576
+    num_attention_heads: 32
+    num_key_value_heads: 2
+    rms_norm_eps: 1.0e-6
+    rope_parameters:
+      factor: 128.0
+      original_max_position_embeddings: 8192
+      rope_theta: 10000
+      rope_type: yarn
+    use_cache: false
   block_size: 8
   num_anchors: 1
   loss_decay_gamma: 7.0
diff --git a/tests/collections/speechlm2/test_salm_dflash.py b/tests/collections/speechlm2/test_salm_dflash.py
index 83a42432d168..259e1609335b 100644
--- a/tests/collections/speechlm2/test_salm_dflash.py
+++ b/tests/collections/speechlm2/test_salm_dflash.py
@@ -12,17 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import Mock
 
 import pytest
 import torch
+from omegaconf import OmegaConf
 from torch import nn
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
 
 from nemo.collections.speechlm2.parts import dflash as salm_dflash
 
 
+REPO_ROOT = Path(__file__).parents[3]
+
+
 class _FakeMoEMesh:
     mesh_dim_names = ("ep_shard", "ep")
 
@@ -142,6 +147,49 @@ def test_build_draft_config_rejects_managed_overrides():
         )
 
 
+def test_salm_automodel_dflash_defaults_match_nemotron_3_5_lightning():
+    cfg = OmegaConf.load(REPO_ROOT / "examples/speechlm2/conf/salm_automodel.yaml")
+    dflash_cfg = OmegaConf.to_container(cfg.dflash, resolve=True)
+    target_config = Qwen3Config(
+        hidden_size=2688,
+        intermediate_size=1856,
+        num_attention_heads=32,
+        num_key_value_heads=2,
+        num_hidden_layers=52,
+        head_dim=128,
+        vocab_size=131072,
+    )
+
+    draft_config, target_layer_ids = salm_dflash._build_draft_config(
+        target_config,
+        dflash_cfg,
+        block_size=dflash_cfg["block_size"],
+        mask_token_id=dflash_cfg["mask_token_id"],
+    )
+
+    assert dflash_cfg["enabled"] is False
+    assert draft_config.num_hidden_layers == 6
+    assert draft_config.hidden_size == 2688
+    assert draft_config.intermediate_size == 6144
+    assert draft_config.num_attention_heads == 32
+    assert draft_config.num_key_value_heads == 2
+    assert draft_config.head_dim == 128
+    assert draft_config.rms_norm_eps == pytest.approx(1.0e-6)
+    assert draft_config.max_position_embeddings == 1048576
+    assert draft_config.rope_parameters == {
+        "factor": 128.0,
+        "original_max_position_embeddings": 8192,
+        "rope_theta": 10000,
+        "rope_type": "yarn",
+    }
+    assert target_layer_ids == [1, 5, 19, 29, 41, 51]
+    assert draft_config.dflash_config == {
+        "mask_token_id": 990,
+        "target_layer_ids": [1, 5, 19, 29, 41, 51],
+    }
+    assert draft_config.block_size == 8
+
+
 class _TargetLLM(nn.Module):
     def __init__(self):
         super().__init__()

From e2183825c9339da1251c52411fc270a617b19a46 Mon Sep 17 00:00:00 2001
From: SlyneD <slyned@nvidia.com>
Date: Tue, 18 Aug 2026 17:44:15 -0700
Subject: [PATCH 3/8] feat(speechlm2): enable memory-bounded DFlash defaults

Signed-off-by: SlyneD <slyned@nvidia.com>
---
 examples/speechlm2/conf/salm_automodel.yaml    | 18 +++++++++++-------
 nemo/collections/speechlm2/parts/dflash.py     | 11 +++++++----
 .../collections/speechlm2/test_salm_dflash.py  |  8 ++++++++
 3 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/examples/speechlm2/conf/salm_automodel.yaml b/examples/speechlm2/conf/salm_automodel.yaml
index 1c2a9713d209..6c7bb1749136 100644
--- a/examples/speechlm2/conf/salm_automodel.yaml
+++ b/examples/speechlm2/conf/salm_automodel.yaml
@@ -299,9 +299,10 @@ exp_manager:
 # inference artifact and cannot directly initialize these trainable BF16
 # nn.Linear weights; these settings reproduce its architecture, not its weights.
 #
-# This initial integration supports BSHD batches with pp_size=cp_size=1. Keep
-# num_anchors small: Automodel's generic DFlash trainer currently materializes
-# draft vocabulary logits.
+# This integration supports BSHD batches with pp_size=cp_size=1. The generic
+# Automodel trainer chunks the LM-head projection and caps the rectangular
+# local-microbatch anchor allocation, so the paper-scale 512-anchor budget does
+# not retain a full anchors-by-block-by-vocabulary logits tensor.
 dflash:
   enabled: false
   mask_token_id: 990
@@ -324,10 +325,13 @@ dflash:
       rope_type: yarn
     use_cache: false
   block_size: 8
-  num_anchors: 1
-  loss_decay_gamma: 7.0
-  attention_backend: sdpa
-  activation_checkpointing: false
+  num_anchors: 512
+  max_total_anchors: 512
+  loss_decay_gamma: 4.0
+  attention_backend: flex_attention
+  activation_checkpointing: true
+  use_fused_linear_ce: true
+  linear_ce_chunk_size: 256
   target_dtype: bfloat16
   lr: 6.0e-4
   output_dir: ./outputs/salm_dflash
diff --git a/nemo/collections/speechlm2/parts/dflash.py b/nemo/collections/speechlm2/parts/dflash.py
index 2bda5ce46c88..089a78a8db9b 100644
--- a/nemo/collections/speechlm2/parts/dflash.py
+++ b/nemo/collections/speechlm2/parts/dflash.py
@@ -200,7 +200,7 @@ def __init__(self, target_model: nn.Module, cfg: dict):
         if mask_token_id is None:
             raise ValueError("dflash.mask_token_id must identify an unused token in the target vocabulary")
         self.mask_token_id = int(mask_token_id)
-        self.attention_backend = str(self.dflash_config.get("attention_backend", "sdpa"))
+        self.attention_backend = str(self.dflash_config.get("attention_backend", "flex_attention"))
         self.output_dir = self.dflash_config.get("output_dir")
         self.learning_rate = float(self.dflash_config.get("lr", 6e-4))
         self.draft_model = None
@@ -256,7 +256,7 @@ def configure_model(self) -> None:
         draft_config._attn_implementation = self.attention_backend
         dtype = next(self.target.llm.parameters()).dtype
         self.draft_model = Qwen3DFlashDraftModel(draft_config).to(self.target.device, dtype=dtype)
-        if self.dflash_config.get("activation_checkpointing", False):
+        if self.dflash_config.get("activation_checkpointing", True):
             self.draft_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
 
         self.trainer_module = DFlashTrainerModule(
@@ -266,10 +266,13 @@ def configure_model(self) -> None:
             mask_token_id=self.mask_token_id,
             block_size=self.block_size,
             attention_backend=self.attention_backend,
-            num_anchors=int(self.dflash_config.get("num_anchors", 1)),
-            loss_decay_gamma=self.dflash_config.get("loss_decay_gamma", 7.0),
+            num_anchors=int(self.dflash_config.get("num_anchors", 512)),
+            max_total_anchors=int(self.dflash_config.get("max_total_anchors", 512)),
+            loss_decay_gamma=self.dflash_config.get("loss_decay_gamma", 4.0),
             loss_type=str(self.dflash_config.get("loss_type", "dflash")),
             prefix_weight_base=float(self.dflash_config.get("prefix_weight_base", 0.9)),
+            use_fused_linear_ce=bool(self.dflash_config.get("use_fused_linear_ce", True)),
+            linear_ce_chunk_size=int(self.dflash_config.get("linear_ce_chunk_size", 256)),
         )
 
         device_mesh = self.device_mesh
diff --git a/tests/collections/speechlm2/test_salm_dflash.py b/tests/collections/speechlm2/test_salm_dflash.py
index 259e1609335b..917d3f61b248 100644
--- a/tests/collections/speechlm2/test_salm_dflash.py
+++ b/tests/collections/speechlm2/test_salm_dflash.py
@@ -168,6 +168,14 @@ def test_salm_automodel_dflash_defaults_match_nemotron_3_5_lightning():
     )
 
     assert dflash_cfg["enabled"] is False
+    assert dflash_cfg["block_size"] == 8
+    assert dflash_cfg["num_anchors"] == 512
+    assert dflash_cfg["max_total_anchors"] == 512
+    assert dflash_cfg["loss_decay_gamma"] == pytest.approx(4.0)
+    assert dflash_cfg["attention_backend"] == "flex_attention"
+    assert dflash_cfg["activation_checkpointing"] is True
+    assert dflash_cfg["use_fused_linear_ce"] is True
+    assert dflash_cfg["linear_ce_chunk_size"] == 256
     assert draft_config.num_hidden_layers == 6
     assert draft_config.hidden_size == 2688
     assert draft_config.intermediate_size == 6144

From 116631f5899316db52a4b0900b21ed722a6ca7a8 Mon Sep 17 00:00:00 2001
From: SlyneD <slyned@nvidia.com>
Date: Fri, 21 Aug 2026 15:48:20 -0700
Subject: [PATCH 4/8] fix(speechlm2): harden SALM DFlash training

Signed-off-by: SlyneD <slyned@nvidia.com>
---
 examples/speechlm2/conf/salm_automodel.yaml   |   2 +-
 examples/speechlm2/salm_train.py              |   5 +-
 nemo/collections/speechlm2/parts/dflash.py    | 145 ++++++++--
 .../collections/speechlm2/test_salm_dflash.py | 272 +++++++++++++++++-
 4 files changed, 398 insertions(+), 26 deletions(-)

diff --git a/examples/speechlm2/conf/salm_automodel.yaml b/examples/speechlm2/conf/salm_automodel.yaml
index 6c7bb1749136..4827bd92cd78 100644
--- a/examples/speechlm2/conf/salm_automodel.yaml
+++ b/examples/speechlm2/conf/salm_automodel.yaml
@@ -299,7 +299,7 @@ exp_manager:
 # inference artifact and cannot directly initialize these trainable BF16
 # nn.Linear weights; these settings reproduce its architecture, not its weights.
 #
-# This integration supports BSHD batches with pp_size=cp_size=1. The generic
+# This integration supports BSHD batches with tp_size=pp_size=cp_size=1. The generic
 # Automodel trainer chunks the LM-head projection and caps the rectangular
 # local-microbatch anchor allocation, so the paper-scale 512-anchor budget does
 # not retain a full anchors-by-block-by-vocabulary logits tensor.
diff --git a/examples/speechlm2/salm_train.py b/examples/speechlm2/salm_train.py
index 1e1607b96567..1bec8930e71f 100644
--- a/examples/speechlm2/salm_train.py
+++ b/examples/speechlm2/salm_train.py
@@ -65,7 +65,10 @@ def train(cfg):
         model = SALMDFlashModule(target_model, OmegaConf.to_container(cfg, resolve=True))
         dataset = _create_salm_dataset(target_model.tokenizer, cfg.data)
         datamodule = DataModule(cfg.data, tokenizer=target_model.tokenizer, dataset=dataset)
-        trainer.fit(model, datamodule)
+        if cfg.get("run_validate_only", False):
+            trainer.validate(model, datamodule)
+        else:
+            trainer.fit(model, datamodule)
         if torch.distributed.is_initialized():
             torch.distributed.destroy_process_group()
         return
diff --git a/nemo/collections/speechlm2/parts/dflash.py b/nemo/collections/speechlm2/parts/dflash.py
index 089a78a8db9b..8b7f5e32e6d6 100644
--- a/nemo/collections/speechlm2/parts/dflash.py
+++ b/nemo/collections/speechlm2/parts/dflash.py
@@ -22,13 +22,24 @@
 
 import torch
 from lightning import LightningModule
-from nemo_automodel.components.speculative.dflash.core import DFlashTrainerModule, NoValidAnchorsError
-from nemo_automodel.components.speculative.dflash.draft_qwen3 import Qwen3DFlashDraftModel, build_target_layer_ids
+from nemo_automodel.components.speculative.dflash.core import (
+    DFlashTrainerModule,
+    NoValidAnchorsError,
+)
+from nemo_automodel.components.speculative.dflash.draft_qwen3 import (
+    Qwen3DFlashDraftModel,
+    build_target_layer_ids,
+)
 from torch import nn
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
 
-from nemo.collections.speechlm2.models.salm import replace_placeholders_and_build_targets
-from nemo.collections.speechlm2.parts.cp_helpers import encode_audio_with_cp_distribution, get_perception_fsdp_group
+from nemo.collections.speechlm2.models.salm import (
+    replace_placeholders_and_build_targets,
+)
+from nemo.collections.speechlm2.parts.cp_helpers import (
+    encode_audio_with_cp_distribution,
+    get_perception_fsdp_group,
+)
 from nemo.core.classes.common import safe_instantiate
 
 _DRAFT_CONFIG_MANAGED_KEYS = {
@@ -94,6 +105,19 @@ def _has_valid_dflash_anchors(loss_mask: torch.Tensor, block_size: int) -> bool:
     return bool((loss_mask[:, : max_anchor + 1] > 0.5).any().item())
 
 
+def _validate_dflash_parallelism(mesh_context) -> None:
+    """Reject parallel layouts whose sequence shards the non-TP DFlash draft cannot consume."""
+    unsupported = []
+    for name in ("tp", "pp", "cp"):
+        size = int(getattr(mesh_context, f"{name}_size", 1))
+        if size > 1:
+            unsupported.append(f"{name}_size={size}")
+    if unsupported:
+        raise NotImplementedError(
+            "SALM DFlash currently requires tp_size=pp_size=cp_size=1; got " + ", ".join(unsupported)
+        )
+
+
 def _build_draft_config(
     target_config,
     dflash_config: dict,
@@ -179,7 +203,10 @@ def _get_consolidated_model_state_dict(model: nn.Module) -> dict[str, torch.Tens
     if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
         return model.state_dict()
 
-    from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
+    from torch.distributed.checkpoint.state_dict import (
+        StateDictOptions,
+        get_model_state_dict,
+    )
 
     return get_model_state_dict(model, options=StateDictOptions(full_state_dict=True, cpu_offload=True))
 
@@ -206,9 +233,17 @@ def __init__(self, target_model: nn.Module, cfg: dict):
         self.draft_model = None
         self.trainer_module = None
         self.target_layer_ids = None
+        self._draft_dp_group = None
+        self._draft_dp_size = 1
         self._partial_val_metrics = defaultdict(list)
         self.register_state_dict_post_hook(self._keep_draft_checkpoint_state)
 
+    def train(self, mode: bool = True):
+        """Set the draft's mode while keeping the frozen target in evaluation mode."""
+        super().train(mode)
+        self.target.eval()
+        return self
+
     @staticmethod
     def _keep_draft_checkpoint_state(module, state_dict, prefix, local_metadata) -> None:
         """Exclude the rebuilt frozen target from resumable checkpoints."""
@@ -236,8 +271,7 @@ def configure_model(self) -> None:
         distributed_setup = getattr(strategy, "distributed_setup", None)
         if distributed_setup is not None:
             mesh_context = distributed_setup.mesh_context
-            if mesh_context.pp_size > 1 or mesh_context.cp_size > 1:
-                raise NotImplementedError("SALM DFlash currently requires pp_size=cp_size=1")
+            _validate_dflash_parallelism(mesh_context)
         self.target._trainer = self.trainer
         self.target.configure_model(
             distributed_setup=distributed_setup,
@@ -283,6 +317,8 @@ def configure_model(self) -> None:
             draft_fsdp_mesh = device_mesh["dp_shard_cp"]
         else:
             draft_fsdp_mesh = device_mesh["dp"]
+        self._draft_dp_size = int(draft_fsdp_mesh.size())
+        self._draft_dp_group = draft_fsdp_mesh.get_group() if self._draft_dp_size > 1 else None
         if draft_fsdp_mesh.size() > 1:
             from torch.distributed.fsdp import fully_shard
 
@@ -340,10 +376,14 @@ def _prepare_batch(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tens
             self.mask_token_id,
         )
         return {
-            "input_ids": expanded_ids[:, :-1],
-            "input_embeddings": input_embeddings[:, :-1],
-            "attention_mask": attention_mask[:, :-1],
-            "loss_mask": target_ids[:, 1:].ne(-100),
+            # DFlash consumes the full, unshifted token stream. Its block builder
+            # gathers a token and its supervision mask at the same sequence index;
+            # applying causal-LM input/label shifting here would offset response
+            # boundaries and remove the final token from draft supervision.
+            "input_ids": expanded_ids,
+            "input_embeddings": input_embeddings,
+            "attention_mask": attention_mask,
+            "loss_mask": target_ids.ne(-100),
         }
 
     @torch.no_grad()
@@ -377,12 +417,19 @@ def hook(_module, _args, output):
         forward_kwargs = {
             "inputs_embeds": inputs["input_embeddings"],
             "attention_mask": inputs["attention_mask"],
+        }
+        forward_parameters = inspect.signature(type(self.target.llm).forward).parameters
+        accepts_extra_kwargs = any(
+            parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in forward_parameters.values()
+        )
+        for name, value in {
             "output_hidden_states": False,
             "use_cache": False,
             "return_dict": True,
-        }
-        if "compute_logits" in inspect.signature(type(self.target.llm).forward).parameters:
-            forward_kwargs["compute_logits"] = False
+            "compute_logits": False,
+        }.items():
+            if accepts_extra_kwargs or name in forward_parameters:
+                forward_kwargs[name] = value
         try:
             self.target.llm(**forward_kwargs)
         finally:
@@ -395,7 +442,8 @@ def hook(_module, _args, output):
     def _run_batch(self, batch: dict[str, torch.Tensor]):
         inputs = self._prepare_batch(batch)
         if not _all_ranks_agree(
-            _has_valid_dflash_anchors(inputs["loss_mask"], self.block_size), inputs["loss_mask"].device
+            _has_valid_dflash_anchors(inputs["loss_mask"], self.block_size),
+            inputs["loss_mask"].device,
         ):
             raise NoValidAnchorsError("At least one rank has no valid DFlash anchors")
         _synchronize_ep_group_before_target_forward(getattr(self.trainer.strategy, "moe_mesh", None))
@@ -406,28 +454,61 @@ def _run_batch(self, batch: dict[str, torch.Tensor]):
             loss_mask=inputs["loss_mask"],
         )
 
+    def _globally_normalized_loss(self, metrics) -> torch.Tensor:
+        """Weight a local DFlash mean by the global draft-DP loss denominator.
+
+        FSDP averages gradients across its process group. Multiplying the local
+        weighted-loss numerator by ``dp_size / global_weight`` therefore yields
+        the true global weighted mean even when ranks sample different numbers of
+        valid anchors or supervised block positions.
+        """
+        local_weight = metrics.loss_weight.to(device=metrics.loss.device, dtype=metrics.loss.dtype)
+        if not (self._draft_dp_size > 1 and torch.distributed.is_available() and torch.distributed.is_initialized()):
+            return metrics.loss
+
+        global_weight = local_weight.detach().clone()
+        torch.distributed.all_reduce(
+            global_weight,
+            op=torch.distributed.ReduceOp.SUM,
+            group=self._draft_dp_group,
+        )
+        return metrics.loss * local_weight * self._draft_dp_size / global_weight.clamp_min(1.0e-6)
+
     def training_step(self, batch, batch_idx):
         batches = list(batch.values()) if isinstance(batch, dict) and "input_ids" not in batch else [batch]
         losses = []
+        skip_counts = defaultdict(int)
         num_batches = _max_rank_value(len(batches), self.device)
         for dataset_index in range(num_batches):
             dataset_batch = batches[dataset_index] if dataset_index < len(batches) else None
             if not _all_ranks_agree(dataset_batch is not None, self.device):
+                skip_counts["missing_batch"] += 1
                 continue
             assert dataset_batch is not None
             signature = _preprocessing_signature(dataset_batch)
             if not _all_ranks_report_same_value(signature, self.device):
+                skip_counts["preprocessing_signature"] += 1
                 continue
             try:
                 metrics = self._run_batch(dataset_batch)
             except NoValidAnchorsError:
+                skip_counts["no_valid_anchors"] += 1
                 continue
-            losses.append(metrics.loss)
+            losses.append(self._globally_normalized_loss(metrics))
             self.log("train/dflash_loss", metrics.loss, on_step=True, prog_bar=True)
             self.log("train/dflash_accuracy", metrics.accuracy, on_step=True)
             self.log("train/accept_len", metrics.accept_len, on_step=True)
         if not losses:
+            # Every rank takes the same synchronized skip branches above. Lightning
+            # rejects ``None`` from ``training_step`` under distributed automatic
+            # optimization, so return a standalone differentiable zero. It has no
+            # graph edge to optimizer-owned draft parameters: backward is valid, all
+            # draft gradients stay ``None``, and AdamW performs no parameter update.
+            self.log("train/dflash_skipped_step", 1.0, on_step=True)
+            for reason, count in skip_counts.items():
+                self.log(f"train/dflash_skip/{reason}", float(count), on_step=True)
             return torch.zeros((), device=self.device, requires_grad=True)
+        self.log("train/dflash_skipped_step", 0.0, on_step=True)
         return torch.stack(losses).mean()
 
     def on_validation_epoch_start(self) -> None:
@@ -450,12 +531,15 @@ def validation_step(self, batch, batch_idx) -> None:
                 metrics = self._run_batch(dataset_batch)
             except NoValidAnchorsError:
                 continue
-            metric_dtype = metrics.loss.dtype
+            # Counts can exceed float32's exact-integer range over a long epoch
+            # with 512 anchors. Accumulate all additive validation statistics in
+            # float64 so a single stacked all-reduce remains exact for counts.
+            metric_dtype = torch.float64
             metric_device = metrics.loss.device
             self._partial_val_metrics[dataset_name].append(
                 torch.stack(
                     [
-                        metrics.loss.detach() * metrics.loss_weight.to(metric_dtype),
+                        metrics.loss.detach().to(dtype=metric_dtype) * metrics.loss_weight.to(metric_dtype),
                         metrics.loss_weight.to(dtype=metric_dtype, device=metric_device),
                         metrics.correct_tokens.to(dtype=metric_dtype, device=metric_device),
                         metrics.valid_tokens.to(dtype=metric_dtype, device=metric_device),
@@ -481,9 +565,22 @@ def on_validation_epoch_end(self) -> None:
 
     def _log_validation_metrics(self, metric_sums: torch.Tensor, suffix: str = "") -> None:
         loss_sum, loss_weight, correct, valid, accept_sum, valid_blocks = metric_sums
-        self.log(f"val/dflash_loss{suffix}", loss_sum / loss_weight.clamp_min(1), on_epoch=True)
-        self.log(f"val/dflash_accuracy{suffix}", correct / valid.clamp_min(1), on_epoch=True)
-        self.log(f"val/accept_len{suffix}", accept_sum / valid_blocks.clamp_min(1), on_epoch=True)
+        accuracy = correct / valid.clamp_min(1)
+        self.log(
+            f"val/dflash_loss{suffix}",
+            loss_sum / loss_weight.clamp_min(1),
+            on_epoch=True,
+        )
+        self.log(f"val/dflash_accuracy{suffix}", accuracy, on_epoch=True)
+        self.log(
+            f"val/accept_len{suffix}",
+            accept_sum / valid_blocks.clamp_min(1),
+            on_epoch=True,
+        )
+        if not suffix:
+            # Preserve the existing SALM recipe's ModelCheckpoint monitor without
+            # changing non-DFlash logging or requiring a DFlash-only exp_manager.
+            self.log("val_acc", accuracy, on_epoch=True)
 
     def configure_optimizers(self):
         optimizer_config = self.dflash_config.get("optimizer")
@@ -500,7 +597,11 @@ def configure_optimizers(self):
         scheduler = safe_instantiate(scheduler_config, optimizer=optimizer, _convert_="all")
         return {
             "optimizer": optimizer,
-            "lr_scheduler": {"scheduler": scheduler, "interval": "step", "frequency": 1},
+            "lr_scheduler": {
+                "scheduler": scheduler,
+                "interval": "step",
+                "frequency": 1,
+            },
         }
 
     def on_train_end(self) -> None:
diff --git a/tests/collections/speechlm2/test_salm_dflash.py b/tests/collections/speechlm2/test_salm_dflash.py
index 917d3f61b248..9f97fc4e86b4 100644
--- a/tests/collections/speechlm2/test_salm_dflash.py
+++ b/tests/collections/speechlm2/test_salm_dflash.py
@@ -22,6 +22,8 @@
 from torch import nn
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
 
+from nemo_automodel.components.loss.dllm_loss import DFlashDecayLoss
+
 from nemo.collections.speechlm2.parts import dflash as salm_dflash
 
 
@@ -66,6 +68,13 @@ def test_synchronize_ep_group_is_noop_without_distributed_ep(monkeypatch):
     assert calls == []
 
 
+def test_validate_dflash_parallelism_rejects_tensor_parallelism():
+    mesh_context = SimpleNamespace(tp_size=2, pp_size=1, cp_size=1)
+
+    with pytest.raises(NotImplementedError, match="tp_size=2"):
+        salm_dflash._validate_dflash_parallelism(mesh_context)
+
+
 def test_expand_ids_with_audio_preserves_internal_pad_valued_tokens():
     input_ids = torch.tensor([[0, 0, 11, 99, 0, 12]])
     audio_embeddings = [torch.randn(3, 4)]
@@ -107,6 +116,78 @@ def test_expand_ids_with_audio_requires_every_replacement_to_be_used():
         )
 
 
+class _BatchTarget(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(1))
+        self.cfg = {}
+        self.text_pad_id = 0
+        self.audio_locator_tag_id = 99
+
+    def _embed_tokens(self, input_ids):
+        return input_ids.to(torch.float32).unsqueeze(-1).expand(-1, -1, 4).clone()
+
+
+class _CaptureDFlashTrainer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.kwargs = None
+
+    def forward(self, **kwargs):
+        self.kwargs = kwargs
+        return "dflash-result"
+
+
+def test_prepare_batch_keeps_full_unshifted_ids_and_token_aligned_loss_mask(
+    monkeypatch,
+):
+    module = salm_dflash.SALMDFlashModule(_BatchTarget(), {"dflash": {"mask_token_id": 990, "block_size": 2}})
+    audio_embeddings = [
+        torch.tensor(
+            [
+                [100.0, 100.0, 100.0, 100.0],
+                [101.0, 101.0, 101.0, 101.0],
+            ]
+        )
+    ]
+    monkeypatch.setattr(module, "_audio_embeddings", Mock(return_value=audio_embeddings))
+    batch = {
+        "input_ids": torch.tensor([[0, 10, 99, 20, 21, 22]]),
+        "loss_mask": torch.tensor([[False, False, False, False, True, True]]),
+    }
+
+    prepared = module._prepare_batch(batch)
+
+    assert prepared["input_ids"].tolist() == [[10, 990, 990, 20, 21, 22]]
+    assert prepared["loss_mask"].tolist() == [[False, False, False, False, True, True]]
+    assert prepared["attention_mask"].tolist() == [[True, True, True, True, True, True]]
+    assert prepared["input_embeddings"].shape == (1, 6, 4)
+    assert prepared["input_embeddings"][0].tolist() == [
+        [10.0, 10.0, 10.0, 10.0],
+        [100.0, 100.0, 100.0, 100.0],
+        [101.0, 101.0, 101.0, 101.0],
+        [20.0, 20.0, 20.0, 20.0],
+        [21.0, 21.0, 21.0, 21.0],
+        [22.0, 22.0, 22.0, 22.0],
+    ]
+
+    captured_hidden = torch.randn(1, 6, 8)
+    target_hidden_states = Mock(return_value=captured_hidden)
+    monkeypatch.setattr(module, "_target_hidden_states", target_hidden_states)
+    module.trainer_module = _CaptureDFlashTrainer()
+    module._trainer = SimpleNamespace(strategy=SimpleNamespace(moe_mesh=None))
+
+    result = module._run_batch(batch)
+
+    assert result == "dflash-result"
+    target_inputs = target_hidden_states.call_args.args[0]
+    assert target_inputs["input_ids"].tolist() == [[10, 990, 990, 20, 21, 22]]
+    assert target_inputs["loss_mask"].tolist() == [[False, False, False, False, True, True]]
+    assert module.trainer_module.kwargs["input_ids"].tolist() == [[10, 990, 990, 20, 21, 22]]
+    assert module.trainer_module.kwargs["loss_mask"].tolist() == [[False, False, False, False, True, True]]
+    assert module.trainer_module.kwargs["hidden_states"] is captured_hidden
+
+
 def test_build_draft_config_applies_explicit_architecture_and_layer_taps():
     target_config = Qwen3Config(
         hidden_size=64,
@@ -132,7 +213,10 @@ def test_build_draft_config_applies_explicit_architecture_and_layer_taps():
     assert layer_ids == [1, 6]
     assert draft_config.num_hidden_layers == 2
     assert draft_config.intermediate_size == 96
-    assert draft_config.dflash_config == {"mask_token_id": 18, "target_layer_ids": [1, 6]}
+    assert draft_config.dflash_config == {
+        "mask_token_id": 18,
+        "target_layer_ids": [1, 6],
+    }
 
 
 def test_build_draft_config_rejects_managed_overrides():
@@ -236,6 +320,23 @@ def __init__(self):
         self.llm = _TargetLLM()
 
 
+class _MinimalTargetLLM(nn.Module):
+    """Target whose explicit forward rejects every optional HF-style kwarg."""
+
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(1))
+        self.layers = nn.ModuleList([nn.Identity(), nn.Identity()])
+        self.calls = []
+
+    def forward(self, *, inputs_embeds, attention_mask):
+        self.calls.append({"attention_mask": attention_mask})
+        hidden = inputs_embeds
+        for index, layer in enumerate(self.layers, start=1):
+            hidden = layer(hidden + index)
+        return hidden
+
+
 def test_target_hidden_states_uses_audio_embeddings_and_skips_logits():
     module = salm_dflash.SALMDFlashModule(_TargetModel(), {"dflash": {"mask_token_id": 18}})
     module.target_layer_ids = [0, 2]
@@ -260,7 +361,25 @@ def test_target_hidden_states_uses_audio_embeddings_and_skips_logits():
     ]
 
 
-def test_get_consolidated_state_dict_uses_plain_state_dict_without_distributed(monkeypatch):
+def test_target_hidden_states_filters_unsupported_optional_forward_kwargs():
+    target = _TargetModel()
+    target.llm = _MinimalTargetLLM()
+    module = salm_dflash.SALMDFlashModule(target, {"dflash": {"mask_token_id": 18}})
+    module.target_layer_ids = [0, 1]
+    inputs = {
+        "input_embeddings": torch.randn(1, 4, 3),
+        "attention_mask": torch.ones(1, 4, dtype=torch.bool),
+    }
+
+    hidden = module._target_hidden_states(inputs)
+
+    assert hidden.shape == (1, 4, 6)
+    assert target.llm.calls == [{"attention_mask": inputs["attention_mask"]}]
+
+
+def test_get_consolidated_state_dict_uses_plain_state_dict_without_distributed(
+    monkeypatch,
+):
     expected = {"weight": torch.tensor([1.0])}
     model = SimpleNamespace(state_dict=Mock(return_value=expected))
     monkeypatch.setattr(salm_dflash.torch.distributed, "is_available", lambda: True)
@@ -272,6 +391,155 @@ def test_get_consolidated_state_dict_uses_plain_state_dict_without_distributed(m
     model.state_dict.assert_called_once_with()
 
 
+def test_train_keeps_frozen_target_in_eval_mode_and_draft_in_requested_mode():
+    target = nn.Sequential(nn.Dropout(p=0.5))
+    module = salm_dflash.SALMDFlashModule(target, {"dflash": {"mask_token_id": 18}})
+    module.draft_model = nn.Sequential(nn.Dropout(p=0.5))
+
+    module.train()
+
+    assert module.training
+    assert module.draft_model.training
+    assert not module.target.training
+    assert not module.target[0].training
+
+
+def test_globally_normalized_loss_uses_draft_dp_weight(monkeypatch):
+    module = salm_dflash.SALMDFlashModule(nn.Linear(1, 1), {"dflash": {"mask_token_id": 18}})
+    module._draft_dp_size = 2
+    module._draft_dp_group = object()
+    monkeypatch.setattr(salm_dflash.torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(salm_dflash.torch.distributed, "is_initialized", lambda: True)
+
+    def fake_all_reduce(value, *, op, group):
+        assert op == salm_dflash.torch.distributed.ReduceOp.SUM
+        assert group is module._draft_dp_group
+        value.fill_(10.0)
+
+    monkeypatch.setattr(salm_dflash.torch.distributed, "all_reduce", fake_all_reduce)
+    local_loss = torch.tensor(2.0, requires_grad=True)
+    metrics = SimpleNamespace(loss=local_loss, loss_weight=torch.tensor(3.0))
+
+    loss = module._globally_normalized_loss(metrics)
+    loss.backward()
+
+    assert loss.item() == pytest.approx(1.2)
+    assert local_loss.grad.item() == pytest.approx(0.6)
+
+
+def test_dflash_loss_times_weight_recovers_decay_weighted_numerator():
+    torch.manual_seed(7)
+    block_size = 4
+    logits = torch.randn(1, 6, 11)
+    targets = torch.randint(0, 11, (1, 6))
+    block_mask = torch.tensor([[1.0, 1.0, 0.0, 1.0, 1.0, 1.0]])
+    loss_fn = DFlashDecayLoss(loss_gamma=4.0, normalize="mean")
+
+    result = loss_fn(logits, targets, block_mask, block_size=block_size)
+
+    nll = torch.nn.functional.cross_entropy(logits.view(-1, 11), targets.view(-1), reduction="none").view(1, 6)
+    depth_weights = torch.exp(-torch.arange(block_size - 1, dtype=logits.dtype) / 4.0).repeat(2)
+    effective_weights = block_mask * depth_weights.unsqueeze(0)
+    expected_numerator = (nll * effective_weights).sum()
+    torch.testing.assert_close(result.total_loss * effective_weights.sum(), expected_numerator)
+
+
+def test_training_step_synchronizes_multi_dataset_skips(monkeypatch):
+    module = salm_dflash.SALMDFlashModule(nn.Linear(1, 1), {"dflash": {"mask_token_id": 18}})
+    module.draft_model = nn.Linear(1, 1)
+    module._draft_dp_size = 1
+    module._draft_dp_group = None
+    monkeypatch.setattr(module, "log", Mock())
+    monkeypatch.setattr(salm_dflash, "_max_rank_value", lambda _value, _device: 3)
+    availability = []
+
+    def agree(local_condition, _device):
+        availability.append(local_condition)
+        return local_condition
+
+    monkeypatch.setattr(salm_dflash, "_all_ranks_agree", agree)
+    monkeypatch.setattr(salm_dflash, "_all_ranks_report_same_value", lambda _value, _device: True)
+    metrics = SimpleNamespace(
+        loss=torch.tensor(2.0, requires_grad=True),
+        loss_weight=torch.tensor(3.0),
+        accuracy=torch.tensor(0.5),
+        accept_len=torch.tensor(1.5),
+    )
+    run_batch = Mock(side_effect=[salm_dflash.NoValidAnchorsError("skip"), metrics])
+    monkeypatch.setattr(module, "_run_batch", run_batch)
+    batch = {
+        "dataset_a": {"input_ids": torch.ones(1, 2, dtype=torch.long)},
+        "dataset_b": {"input_ids": torch.ones(1, 2, dtype=torch.long)},
+    }
+
+    loss = module.training_step(batch, batch_idx=0)
+
+    torch.testing.assert_close(loss, metrics.loss)
+    assert availability == [True, True, False]
+    assert run_batch.call_count == 2
+
+
+def test_training_step_returns_differentiable_zero_when_every_dataset_is_skipped(monkeypatch):
+    module = salm_dflash.SALMDFlashModule(nn.Linear(1, 1), {"dflash": {"mask_token_id": 18}})
+    module.draft_model = nn.Linear(1, 1)
+    log = Mock()
+    monkeypatch.setattr(module, "log", log)
+    monkeypatch.setattr(salm_dflash, "_max_rank_value", lambda value, _device: value)
+    monkeypatch.setattr(salm_dflash, "_all_ranks_agree", lambda condition, _device: condition)
+    monkeypatch.setattr(salm_dflash, "_all_ranks_report_same_value", lambda _value, _device: True)
+    monkeypatch.setattr(
+        module,
+        "_run_batch",
+        Mock(side_effect=salm_dflash.NoValidAnchorsError("skip")),
+    )
+
+    loss = module.training_step({"input_ids": torch.ones(1, 2, dtype=torch.long)}, batch_idx=0)
+
+    assert loss.item() == 0.0
+    assert loss.requires_grad
+    loss.backward()
+    assert all(parameter.grad is None for parameter in module.draft_model.parameters())
+    log.assert_any_call("train/dflash_skipped_step", 1.0, on_step=True)
+    log.assert_any_call("train/dflash_skip/no_valid_anchors", 1.0, on_step=True)
+
+
+def test_validation_step_accumulates_additive_metrics_in_float64(monkeypatch):
+    module = salm_dflash.SALMDFlashModule(nn.Linear(1, 1), {"dflash": {"mask_token_id": 18}})
+    module.draft_model = nn.Linear(1, 1)
+    metrics = SimpleNamespace(
+        loss=torch.tensor(2.0, dtype=torch.bfloat16),
+        loss_weight=torch.tensor(3.0),
+        correct_tokens=torch.tensor(2**24 + 1),
+        valid_tokens=torch.tensor(2**24 + 3),
+        accept_len_sum=torch.tensor(7.0),
+        valid_blocks=torch.tensor(4),
+    )
+    monkeypatch.setattr(salm_dflash, "_max_rank_value", lambda value, _device: value)
+    monkeypatch.setattr(salm_dflash, "_all_ranks_agree", lambda condition, _device: condition)
+    monkeypatch.setattr(salm_dflash, "_all_ranks_report_same_value", lambda _value, _device: True)
+    monkeypatch.setattr(module, "_run_batch", Mock(return_value=metrics))
+
+    module.validation_step({"input_ids": torch.ones(1, 2, dtype=torch.long)}, batch_idx=0)
+
+    stored = module._partial_val_metrics["validation"][0]
+    assert stored.dtype == torch.float64
+    assert stored[2].item() == 2**24 + 1
+    assert stored[3].item() == 2**24 + 3
+
+
+def test_aggregate_validation_accuracy_preserves_default_checkpoint_monitor(
+    monkeypatch,
+):
+    module = salm_dflash.SALMDFlashModule(nn.Linear(1, 1), {"dflash": {"mask_token_id": 18}})
+    log = Mock()
+    monkeypatch.setattr(module, "log", log)
+
+    module._log_validation_metrics(torch.tensor([8.0, 4.0, 3.0, 6.0, 5.0, 2.0]))
+
+    log.assert_any_call("val/dflash_accuracy", torch.tensor(0.5), on_epoch=True)
+    log.assert_any_call("val_acc", torch.tensor(0.5), on_epoch=True)
+
+
 def test_state_dict_hook_keeps_only_draft_parameters():
     module = SimpleNamespace(_CHECKPOINT_STATE_PREFIX="draft_model.")
     state_dict = {

From 54e087b37fb8fdd2150d96f5d34ca368a0f7f1ef Mon Sep 17 00:00:00 2001
From: SlyneD <slyned@nvidia.com>
Date: Tue, 25 Aug 2026 14:00:54 -0700
Subject: [PATCH 5/8] fix(speechlm2): finalize SALM DFlash integration

Signed-off-by: SlyneD <slyned@nvidia.com>
---
 docs/source/speechlm2/intro.rst               | 11 ++++
 nemo/collections/speechlm2/parts/dflash.py    | 26 +++++++--
 pyproject.toml                                |  6 ++-
 .../collections/speechlm2/test_salm_dflash.py | 54 ++++++++++++++++++-
 uv.lock                                       | 12 ++---
 5 files changed, 95 insertions(+), 14 deletions(-)

diff --git a/docs/source/speechlm2/intro.rst b/docs/source/speechlm2/intro.rst
index 8aac63f9f12f..a62e797cf99e 100644
--- a/docs/source/speechlm2/intro.rst
+++ b/docs/source/speechlm2/intro.rst
@@ -370,6 +370,17 @@ The ``salm_automodel.yaml`` config sets ``model.use_nemo_automodel: true``, whic
 ``SALMAutomodel`` class. This variant supports ``AutomodelParallelStrategy`` for FSDP2/TP/EP
 parallelism and MoE optimizations (Grouped GEMM, DeepEP).
 
+DFlash draft training
+~~~~~~~~~~~~~~~~~~~~~
+
+The same config includes an optional ``dflash:`` section for training a compact DFlash draft
+against a frozen, audio-conditioned ``SALMAutomodel`` target. Set ``dflash.enabled=true`` and
+provide a reserved ``dflash.mask_token_id``; ``salm_train.py`` then trains and exports draft-only
+weights while preserving SALM's audio-placeholder expansion. The default anchor budget and fused
+linear cross-entropy bound peak vocabulary-logit memory. This initial integration supports BSHD
+batches with ``tp_size=pp_size=cp_size=1`` and does not directly load the published packed NVFP4
+inference checkpoint into BF16 training modules.
+
 For more detailed information on training at scale, model parallelism, and SLURM-based training, see :doc:`training and scaling <training_and_scaling>`.
 
 Collection Structure
diff --git a/nemo/collections/speechlm2/parts/dflash.py b/nemo/collections/speechlm2/parts/dflash.py
index 8b7f5e32e6d6..99d9a1567819 100644
--- a/nemo/collections/speechlm2/parts/dflash.py
+++ b/nemo/collections/speechlm2/parts/dflash.py
@@ -31,6 +31,7 @@
     build_target_layer_ids,
 )
 from torch import nn
+from torch.distributed.tensor import DTensor
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
 
 from nemo.collections.speechlm2.models.salm import (
@@ -100,7 +101,13 @@ def _synchronize_ep_group_before_target_forward(moe_mesh) -> None:
 
 
 def _has_valid_dflash_anchors(loss_mask: torch.Tensor, block_size: int) -> bool:
-    """Return whether this rank can form at least one DFlash anchor."""
+    """Mirror Automodel's unpacked DFlash anchor-validity predicate.
+
+    SALM DFlash rejects packed sequences, so ``DFlashTrainerModule`` considers
+    an anchor valid exactly when its own position is supervised and it lies no
+    later than ``seq_len - block_size``. Following block positions may be masked;
+    they affect the loss denominator but not anchor validity.
+    """
     max_anchor = max(loss_mask.shape[1] - block_size, 0)
     return bool((loss_mask[:, : max_anchor + 1] > 0.5).any().item())
 
@@ -177,7 +184,10 @@ def _expand_ids_with_audio(
     replacement_idx = 0
     for row in input_ids:
         non_padding = (row != padding_id).nonzero(as_tuple=False)
-        first_non_padding = int(non_padding[0]) if non_padding.numel() else row.numel()
+        # Match input_utils._unpad_inputs exactly: an all-padding row retains its
+        # last element rather than becoming empty, keeping ids and embeddings
+        # aligned even for this degenerate input.
+        first_non_padding = int(non_padding[0]) if non_padding.numel() else row.numel() - 1
         row = row[first_non_padding:]
         pieces = []
         for token in row:
@@ -328,8 +338,13 @@ def configure_model(self) -> None:
             raise RuntimeError("The DFlash SALM target must be fully frozen")
 
     @property
-    def device(self):
-        return next(self.draft_model.parameters()).device
+    def device(self) -> torch.device:
+        """Infer the device from regular or FSDP2-sharded draft parameters."""
+        if self.draft_model is not None:
+            parameter = next(self.draft_model.parameters(), None)
+            if parameter is not None:
+                return parameter._local_tensor.device if isinstance(parameter, DTensor) else parameter.device
+        return super().device
 
     def _audio_embeddings(self, batch: dict[str, torch.Tensor]) -> list[torch.Tensor]:
         spk_targets = batch.get("spk_targets")
@@ -492,6 +507,8 @@ def training_step(self, batch, batch_idx):
             try:
                 metrics = self._run_batch(dataset_batch)
             except NoValidAnchorsError:
+                # _run_batch's exact, synchronized precheck makes this branch
+                # rank-symmetric; retain the catch for direct/test callers.
                 skip_counts["no_valid_anchors"] += 1
                 continue
             losses.append(self._globally_normalized_loss(metrics))
@@ -530,6 +547,7 @@ def validation_step(self, batch, batch_idx) -> None:
             try:
                 metrics = self._run_batch(dataset_batch)
             except NoValidAnchorsError:
+                # Rank symmetry is guaranteed by _run_batch's exact precheck.
                 continue
             # Counts can exceed float32's exact-integer range over a long epoch
             # with 512 anchors. Accumulate all additive validation statistics in
diff --git a/pyproject.toml b/pyproject.toml
index 1341cedafbf4..e48b9c6657db 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -472,8 +472,10 @@ deep-ep = ["setuptools"]
 # --- uv configuration ---
 # Keep Torch wheel indexes explicit per CUDA extra. Automodel's upstream source
 # metadata is mirrored below so this repository controls the CUDA wheel index.
+# Draft-only stacked pin to Automodel PR #3673; replace with its upstream merge
+# commit before marking the dependent NeMo PR ready for review.
 [tool.uv.sources]
-nemo_automodel = { git = "https://github.com/NVIDIA-NeMo/Automodel.git", rev = "7b15c9bfbc620bec4dd3b819fce2e48726a7ddeb" }
+nemo_automodel = { git = "https://github.com/NVIDIA-NeMo/Automodel.git", rev = "refs/pull/3673/head" }
 megatron-fsdp = { git = "https://github.com/yuhezhang-ai/Megatron-LM.git", rev = "455389c480af6b3acdca74c7830c68b3274eb083", subdirectory = "megatron/core/distributed/fsdp/src" }
 deep_ep = { git = "https://github.com/deepseek-ai/DeepEP.git", rev = "7febc6e25660af0f54d95dd781ecdcd62265ecca" }
 torch = [
@@ -507,7 +509,7 @@ explicit = true
 
 [[tool.uv.dependency-metadata]]
 name = "nemo-automodel"
-version = "0.5.0+cb03cf61"
+version = "0.5.0+426711b8"
 requires-python = ">=3.10"
 # Keep Automodel main's core dependency metadata static here so this repo
 # controls the CUDA wheel index.
diff --git a/tests/collections/speechlm2/test_salm_dflash.py b/tests/collections/speechlm2/test_salm_dflash.py
index 9f97fc4e86b4..ff989ec6572c 100644
--- a/tests/collections/speechlm2/test_salm_dflash.py
+++ b/tests/collections/speechlm2/test_salm_dflash.py
@@ -22,14 +22,45 @@
 from torch import nn
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
 
-from nemo_automodel.components.loss.dllm_loss import DFlashDecayLoss
+pytest.importorskip("nemo_automodel")
+pytestmark = pytest.mark.unit
 
-from nemo.collections.speechlm2.parts import dflash as salm_dflash
+from nemo_automodel.components.loss.dllm_loss import DFlashDecayLoss  # noqa: E402
+
+from nemo.collections.speechlm2.parts import dflash as salm_dflash  # noqa: E402
 
 
 REPO_ROOT = Path(__file__).parents[3]
 
 
+@pytest.mark.parametrize(
+    "loss_mask,block_size",
+    [
+        (torch.tensor([[1.0] * 8]), 8),
+        (torch.tensor([[1.0] * 4]), 8),
+        (torch.tensor([[0.0] * 8 + [1.0] * 8]), 8),
+        (torch.tensor([[0.0] * 9 + [1.0] * 7]), 8),
+        (torch.tensor([[0.0] * 16, [0.0] * 8 + [1.0] * 8]), 8),
+    ],
+)
+def test_anchor_precheck_matches_automodel_unpacked_sampler(loss_mask, block_size):
+    """The synchronized precheck must exactly predict Automodel's early raise."""
+    trainer = SimpleNamespace(block_size=block_size, num_anchors=512, max_total_anchors=None)
+    try:
+        salm_dflash.DFlashTrainerModule._sample_anchor_positions(
+            trainer,
+            seq_len=loss_mask.shape[1],
+            loss_mask=loss_mask,
+            device=loss_mask.device,
+        )
+    except salm_dflash.NoValidAnchorsError:
+        automodel_has_valid = False
+    else:
+        automodel_has_valid = True
+
+    assert salm_dflash._has_valid_dflash_anchors(loss_mask, block_size) is automodel_has_valid
+
+
 class _FakeMoEMesh:
     mesh_dim_names = ("ep_shard", "ep")
 
@@ -116,6 +147,25 @@ def test_expand_ids_with_audio_requires_every_replacement_to_be_used():
         )
 
 
+def test_expand_ids_with_audio_matches_unpad_behavior_for_all_padding_row():
+    expanded = salm_dflash._expand_ids_with_audio(
+        torch.tensor([[0, 0, 0], [0, 11, 12]]),
+        [],
+        padding_id=0,
+        placeholder_id=99,
+        mask_token_id=18,
+    )
+
+    assert expanded.tolist() == [[0, 0], [11, 12]]
+
+
+def test_device_falls_back_before_draft_configuration():
+    module = salm_dflash.SALMDFlashModule(nn.Linear(1, 1), {"dflash": {"mask_token_id": 18}})
+
+    assert module.draft_model is None
+    assert module.device == torch.device("cpu")
+
+
 class _BatchTarget(nn.Module):
     def __init__(self):
         super().__init__()
diff --git a/uv.lock b/uv.lock
index 9a7876c3c246..b2d26a681d4a 100644
--- a/uv.lock
+++ b/uv.lock
@@ -58,7 +58,7 @@ overrides = [
 
 [[manifest.dependency-metadata]]
 name = "nemo-automodel"
-version = "0.5.0+cb03cf61"
+version = "0.5.0+426711b8"
 requires-dist = ["datasets>=4.0.0", "megatron-fsdp==0.5.0", "mistral-common[audio,hf-hub,sentencepiece]", "pybind11", "pyyaml", "tiktoken", "torch>=2.6.0", "torchdata", "transformers==5.12.1", "wandb>=0.28.0", "torchao", "mlflow", "flashoptim>=0.1.3", "quack-kernels==0.6.1 ; sys_platform == 'linux'"]
 requires-python = ">=3.10"
 
@@ -4179,8 +4179,8 @@ wheels = [
 
 [[package]]
 name = "nemo-automodel"
-version = "0.5.0+cb03cf61"
-source = { git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=7b15c9bfbc620bec4dd3b819fce2e48726a7ddeb#7b15c9bfbc620bec4dd3b819fce2e48726a7ddeb" }
+version = "0.5.0+426711b8"
+source = { git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=refs%2Fpull%2F3673%2Fhead#426711b8336d563c37c7ec4a9b97b1bb6dd86aca" }
 dependencies = [
     { name = "datasets" },
     { name = "flashoptim" },
@@ -4633,9 +4633,9 @@ requires-dist = [
     { name = "matplotlib", marker = "extra == 'audio'" },
     { name = "matplotlib", marker = "extra == 'speechlm2'" },
     { name = "matplotlib", marker = "extra == 'tts'" },
-    { name = "nemo-automodel", marker = "extra == 'all'", git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=7b15c9bfbc620bec4dd3b819fce2e48726a7ddeb" },
-    { name = "nemo-automodel", marker = "extra == 'speechlm2'", git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=7b15c9bfbc620bec4dd3b819fce2e48726a7ddeb" },
-    { name = "nemo-automodel", marker = "extra == 'speechlm2-only'", git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=7b15c9bfbc620bec4dd3b819fce2e48726a7ddeb" },
+    { name = "nemo-automodel", marker = "extra == 'all'", git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=refs%2Fpull%2F3673%2Fhead" },
+    { name = "nemo-automodel", marker = "extra == 'speechlm2'", git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=refs%2Fpull%2F3673%2Fhead" },
+    { name = "nemo-automodel", marker = "extra == 'speechlm2-only'", git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=refs%2Fpull%2F3673%2Fhead" },
     { name = "nemo-text-processing", marker = "'aarch' not in platform_machine and 'arm' not in platform_machine and sys_platform != 'darwin' and extra == 'all'" },
     { name = "nemo-text-processing", marker = "'aarch' not in platform_machine and 'arm' not in platform_machine and sys_platform != 'darwin' and extra == 'speechlm2'" },
     { name = "nemo-text-processing", marker = "'aarch' not in platform_machine and 'arm' not in platform_machine and sys_platform != 'darwin' and extra == 'tts'" },

From 7ec0e5caebed4ec37b6a072255936be53c891832 Mon Sep 17 00:00:00 2001
From: SlyneD <slyned@nvidia.com>
Date: Tue, 25 Aug 2026 17:26:58 -0700
Subject: [PATCH 6/8] feat(speechlm2): support DFlash2 SALM training

Signed-off-by: SlyneD <slyned@nvidia.com>
---
 docs/source/speechlm2/intro.rst               |  25 +-
 examples/speechlm2/conf/salm_automodel.yaml   |  19 +-
 nemo/collections/speechlm2/parts/dflash.py    | 257 ++++++++++++++----
 pyproject.toml                                |   6 +-
 .../collections/speechlm2/test_salm_dflash.py | 247 ++++++++++++++++-
 uv.lock                                       |   8 +-
 6 files changed, 477 insertions(+), 85 deletions(-)

diff --git a/docs/source/speechlm2/intro.rst b/docs/source/speechlm2/intro.rst
index a62e797cf99e..ce2eb05f1805 100644
--- a/docs/source/speechlm2/intro.rst
+++ b/docs/source/speechlm2/intro.rst
@@ -370,16 +370,21 @@ The ``salm_automodel.yaml`` config sets ``model.use_nemo_automodel: true``, whic
 ``SALMAutomodel`` class. This variant supports ``AutomodelParallelStrategy`` for FSDP2/TP/EP
 parallelism and MoE optimizations (Grouped GEMM, DeepEP).
 
-DFlash draft training
-~~~~~~~~~~~~~~~~~~~~~
-
-The same config includes an optional ``dflash:`` section for training a compact DFlash draft
-against a frozen, audio-conditioned ``SALMAutomodel`` target. Set ``dflash.enabled=true`` and
-provide a reserved ``dflash.mask_token_id``; ``salm_train.py`` then trains and exports draft-only
-weights while preserving SALM's audio-placeholder expansion. The default anchor budget and fused
-linear cross-entropy bound peak vocabulary-logit memory. This initial integration supports BSHD
-batches with ``tp_size=pp_size=cp_size=1`` and does not directly load the published packed NVFP4
-inference checkpoint into BF16 training modules.
+DFlash and DFlash2 draft training
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The same config includes an optional ``dflash:`` section for training a compact DFlash or DFlash2
+draft against a frozen, audio-conditioned ``SALMAutomodel`` target. Set ``dflash.enabled=true``,
+choose ``dflash.variant`` (``dflash`` or ``dflash2``), and provide a reserved
+``dflash.mask_token_id``; ``salm_train.py`` then trains and exports draft-only weights while
+preserving SALM's audio-placeholder expansion. The shipped DFlash2 settings mirror Automodel's
+recipe, including its two-tap grouped dynamic convolution, top-16 rank-256 path selector, and
+separately normalized backbone and selector losses. ``max_total_anchors`` bounds both variants'
+anchor allocation. Fused linear cross-entropy further bounds DFlash vocabulary-logit memory, but
+DFlash2 requires dense logits for candidate selection and therefore requires
+``use_fused_linear_ce=false``. This integration supports BSHD batches with
+``tp_size=pp_size=cp_size=1`` and does not directly load the published packed NVFP4 inference
+checkpoint into BF16 training modules.
 
 For more detailed information on training at scale, model parallelism, and SLURM-based training, see :doc:`training and scaling <training_and_scaling>`.
 
diff --git a/examples/speechlm2/conf/salm_automodel.yaml b/examples/speechlm2/conf/salm_automodel.yaml
index 4827bd92cd78..6e65ed1cae9f 100644
--- a/examples/speechlm2/conf/salm_automodel.yaml
+++ b/examples/speechlm2/conf/salm_automodel.yaml
@@ -299,12 +299,14 @@ exp_manager:
 # inference artifact and cannot directly initialize these trainable BF16
 # nn.Linear weights; these settings reproduce its architecture, not its weights.
 #
-# This integration supports BSHD batches with tp_size=pp_size=cp_size=1. The generic
-# Automodel trainer chunks the LM-head projection and caps the rectangular
-# local-microbatch anchor allocation, so the paper-scale 512-anchor budget does
-# not retain a full anchors-by-block-by-vocabulary logits tensor.
+# This integration supports DFlash and DFlash2 with BSHD batches and
+# tp_size=pp_size=cp_size=1. The DFlash2 defaults mirror Automodel's recipe: a
+# two-tap grouped dynamic convolution and a rank-256 top-16 path selector. DFlash2
+# needs dense draft logits for that selector, so fused linear CE is only available
+# when variant=dflash.
 dflash:
   enabled: false
+  variant: dflash2  # dflash or dflash2
   mask_token_id: 990
   draft_num_hidden_layers: 6
   target_layer_ids: [1, 5, 19, 29, 41, 51]
@@ -328,9 +330,16 @@ dflash:
   num_anchors: 512
   max_total_anchors: 512
   loss_decay_gamma: 4.0
+  draft_sliding_window: null
   attention_backend: flex_attention
   activation_checkpointing: true
-  use_fused_linear_ce: true
+  # DFlash2 in-block convolution and pairwise path selector (Automodel recipe defaults).
+  conv_kernel_size: 2
+  conv_group_size: 16
+  selector_rank: 256
+  selector_top_k: 16
+  selector_loss_weight: 1.0
+  use_fused_linear_ce: false  # set true only with variant=dflash
   linear_ce_chunk_size: 256
   target_dtype: bfloat16
   lr: 6.0e-4
diff --git a/nemo/collections/speechlm2/parts/dflash.py b/nemo/collections/speechlm2/parts/dflash.py
index 99d9a1567819..8bd02ec8977a 100644
--- a/nemo/collections/speechlm2/parts/dflash.py
+++ b/nemo/collections/speechlm2/parts/dflash.py
@@ -22,14 +22,17 @@
 
 import torch
 from lightning import LightningModule
+from nemo_automodel.components.distributed.mesh_utils import get_flat_mesh, get_fsdp_dp_mesh
 from nemo_automodel.components.speculative.dflash.core import (
     DFlashTrainerModule,
     NoValidAnchorsError,
 )
+from nemo_automodel.components.speculative.dflash.dflash2_core import DFlash2TrainerModule
 from nemo_automodel.components.speculative.dflash.draft_qwen3 import (
     Qwen3DFlashDraftModel,
     build_target_layer_ids,
 )
+from nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2 import Qwen3DFlash2DraftModel
 from torch import nn
 from torch.distributed.tensor import DTensor
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
@@ -47,11 +50,15 @@
     "architectures",
     "block_size",
     "dflash_config",
+    "is_causal",
     "layer_types",
     "max_window_layers",
     "num_hidden_layers",
     "num_target_layers",
+    "sliding_window",
+    "use_sliding_window",
 }
+_DFLASH_VARIANTS = {"dflash", "dflash2"}
 
 
 def _all_ranks_agree(local_condition: bool, device: torch.device) -> bool:
@@ -132,6 +139,9 @@ def _build_draft_config(
     mask_token_id: int,
 ) -> tuple[Qwen3Config, list[int]]:
     """Create the Qwen3-shaped DFlash draft config for a SALM target."""
+    variant = str(dflash_config.get("variant", "dflash")).lower()
+    if variant not in _DFLASH_VARIANTS:
+        raise ValueError(f"dflash.variant must be one of {sorted(_DFLASH_VARIANTS)}, got {variant!r}")
     num_target_layers = int(target_config.num_hidden_layers)
     draft_layers = int(dflash_config.get("draft_num_hidden_layers", 2))
     target_layer_ids = list(
@@ -147,20 +157,49 @@ def _build_draft_config(
     if managed:
         raise ValueError(f"dflash.draft_model_config cannot override managed keys: {', '.join(managed)}")
 
+    draft_metadata = {
+        "block_size": block_size,
+        "mask_token_id": mask_token_id,
+        "target_layer_ids": target_layer_ids,
+    }
+    if variant == "dflash2":
+        draft_metadata.update(
+            {
+                "conv_kernel_size": int(dflash_config.get("conv_kernel_size", 2)),
+                "conv_group_size": int(dflash_config.get("conv_group_size", 16)),
+                "selector_rank": int(dflash_config.get("selector_rank", 256)),
+                "selector_top_k": int(dflash_config.get("selector_top_k", 16)),
+            }
+        )
+
+    draft_layers_config = {
+        "layer_types": ["full_attention"] * draft_layers,
+        "sliding_window": None,
+        "use_sliding_window": False,
+    }
+    sliding_window = dflash_config.get("draft_sliding_window")
+    if sliding_window is not None:
+        sliding_window = int(sliding_window)
+        if sliding_window <= 0:
+            raise ValueError(f"dflash.draft_sliding_window must be > 0 or null, got {sliding_window}")
+        draft_layers_config = {
+            "layer_types": ["sliding_attention"] * draft_layers,
+            "sliding_window": sliding_window,
+            "use_sliding_window": True,
+        }
+
     draft_dict = target_config.to_dict()
     draft_dict.update(architecture)
     draft_dict.update(
         {
-            "architectures": ["Qwen3DFlashDraftModel"],
+            "architectures": ["Qwen3DFlash2DraftModel" if variant == "dflash2" else "Qwen3DFlashDraftModel"],
             "num_hidden_layers": draft_layers,
-            "layer_types": ["full_attention"] * draft_layers,
             "max_window_layers": draft_layers,
+            "is_causal": False,
             "num_target_layers": num_target_layers,
             "block_size": block_size,
-            "dflash_config": {
-                "mask_token_id": mask_token_id,
-                "target_layer_ids": target_layer_ids,
-            },
+            "dflash_config": draft_metadata,
+            **draft_layers_config,
         }
     )
     draft_config = Qwen3Config.from_dict(draft_dict)
@@ -232,6 +271,9 @@ def __init__(self, target_model: nn.Module, cfg: dict):
         self.target = target_model
         self.cfg = cfg
         self.dflash_config = cfg.get("dflash", cfg)
+        self.dflash_variant = str(self.dflash_config.get("variant", "dflash")).lower()
+        if self.dflash_variant not in _DFLASH_VARIANTS:
+            raise ValueError(f"dflash.variant must be one of {sorted(_DFLASH_VARIANTS)}, got {self.dflash_variant!r}")
         self.block_size = int(self.dflash_config.get("block_size", 8))
         mask_token_id = self.dflash_config.get("mask_token_id")
         if mask_token_id is None:
@@ -240,6 +282,18 @@ def __init__(self, target_model: nn.Module, cfg: dict):
         self.attention_backend = str(self.dflash_config.get("attention_backend", "flex_attention"))
         self.output_dir = self.dflash_config.get("output_dir")
         self.learning_rate = float(self.dflash_config.get("lr", 6e-4))
+        self.selector_loss_weight = float(self.dflash_config.get("selector_loss_weight", 1.0))
+        if self.selector_loss_weight < 0:
+            raise ValueError(f"dflash.selector_loss_weight must be >= 0, got {self.selector_loss_weight}")
+        if self.dflash_variant == "dflash2":
+            loss_type = str(self.dflash_config.get("loss_type", None) or "dflash")
+            if loss_type != "dflash":
+                raise ValueError("dflash.loss_type must be 'dflash' when dflash.variant='dflash2'")
+            if bool(self.dflash_config.get("use_fused_linear_ce", False)):
+                raise ValueError(
+                    "dflash.use_fused_linear_ce is not supported by DFlash2 because its path selector needs "
+                    "the draft logits; set it to false"
+                )
         self.draft_model = None
         self.trainer_module = None
         self.target_layer_ids = None
@@ -299,36 +353,18 @@ def configure_model(self) -> None:
         )
         draft_config._attn_implementation = self.attention_backend
         dtype = next(self.target.llm.parameters()).dtype
-        self.draft_model = Qwen3DFlashDraftModel(draft_config).to(self.target.device, dtype=dtype)
+        draft_cls = self._draft_model_class()
+        self.draft_model = draft_cls(draft_config).to(self.target.device, dtype=dtype)
         if self.dflash_config.get("activation_checkpointing", True):
             self.draft_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
 
-        self.trainer_module = DFlashTrainerModule(
-            draft_model=self.draft_model,
-            target_lm_head=self.target.llm.get_output_embeddings(),
-            target_embed_tokens=self.target.llm.get_input_embeddings(),
-            mask_token_id=self.mask_token_id,
-            block_size=self.block_size,
-            attention_backend=self.attention_backend,
-            num_anchors=int(self.dflash_config.get("num_anchors", 512)),
-            max_total_anchors=int(self.dflash_config.get("max_total_anchors", 512)),
-            loss_decay_gamma=self.dflash_config.get("loss_decay_gamma", 4.0),
-            loss_type=str(self.dflash_config.get("loss_type", "dflash")),
-            prefix_weight_base=float(self.dflash_config.get("prefix_weight_base", 0.9)),
-            use_fused_linear_ce=bool(self.dflash_config.get("use_fused_linear_ce", True)),
-            linear_ce_chunk_size=int(self.dflash_config.get("linear_ce_chunk_size", 256)),
-        )
+        self.trainer_module = self._create_trainer_module()
 
         device_mesh = self.device_mesh
-        dim_names = device_mesh.mesh_dim_names
-        if "dp_replicate" in dim_names and "dp_shard_cp" in dim_names:
-            draft_fsdp_mesh = device_mesh["dp_replicate", "dp_shard_cp"]
-        elif "dp_shard_cp" in dim_names:
-            draft_fsdp_mesh = device_mesh["dp_shard_cp"]
-        else:
-            draft_fsdp_mesh = device_mesh["dp"]
-        self._draft_dp_size = int(draft_fsdp_mesh.size())
-        self._draft_dp_group = draft_fsdp_mesh.get_group() if self._draft_dp_size > 1 else None
+        draft_fsdp_mesh = get_fsdp_dp_mesh(device_mesh)
+        draft_dp_mesh = get_flat_mesh(device_mesh, "dp")
+        self._draft_dp_size = int(draft_dp_mesh.size())
+        self._draft_dp_group = draft_dp_mesh.get_group() if self._draft_dp_size > 1 else None
         if draft_fsdp_mesh.size() > 1:
             from torch.distributed.fsdp import fully_shard
 
@@ -337,6 +373,38 @@ def configure_model(self) -> None:
         if any(parameter.requires_grad for parameter in self.target.parameters()):
             raise RuntimeError("The DFlash SALM target must be fully frozen")
 
+    def _create_trainer_module(self) -> DFlashTrainerModule:
+        """Build the Automodel trainer matching the configured draft variant."""
+        max_total_anchors = self.dflash_config.get("max_total_anchors", 512)
+        common_trainer_kwargs = {
+            "draft_model": self.draft_model,
+            "target_lm_head": self.target.llm.get_output_embeddings(),
+            "target_embed_tokens": self.target.llm.get_input_embeddings(),
+            "mask_token_id": self.mask_token_id,
+            "block_size": self.block_size,
+            "attention_backend": self.attention_backend,
+            "num_anchors": int(self.dflash_config.get("num_anchors", 512)),
+            "max_total_anchors": int(max_total_anchors) if max_total_anchors is not None else None,
+            "loss_decay_gamma": self.dflash_config.get("loss_decay_gamma", 4.0),
+            "sliding_window": self.dflash_config.get("draft_sliding_window"),
+        }
+        if self.dflash_variant == "dflash2":
+            return DFlash2TrainerModule(
+                **common_trainer_kwargs,
+                selector_loss_weight=self.selector_loss_weight,
+            )
+        return DFlashTrainerModule(
+            **common_trainer_kwargs,
+            loss_type=str(self.dflash_config.get("loss_type", None) or "dflash"),
+            prefix_weight_base=float(self.dflash_config.get("prefix_weight_base", 0.9)),
+            use_fused_linear_ce=bool(self.dflash_config.get("use_fused_linear_ce", True)),
+            linear_ce_chunk_size=int(self.dflash_config.get("linear_ce_chunk_size", 256)),
+        )
+
+    def _draft_model_class(self) -> type[Qwen3DFlashDraftModel]:
+        """Return the draft implementation selected by ``dflash.variant``."""
+        return Qwen3DFlash2DraftModel if self.dflash_variant == "dflash2" else Qwen3DFlashDraftModel
+
     @property
     def device(self) -> torch.device:
         """Infer the device from regular or FSDP2-sharded draft parameters."""
@@ -406,10 +474,13 @@ def _target_hidden_states(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor
         """Run the frozen audio-conditioned target and concatenate configured layers."""
         if hasattr(self.target.llm, "model") and hasattr(self.target.llm.model, "layers"):
             layer_container = self.target.llm.model.layers
+            final_norm = getattr(self.target.llm.model, "norm", None)
         elif hasattr(self.target.llm, "layers"):
             layer_container = self.target.llm.layers
+            final_norm = getattr(self.target.llm, "norm", None)
         elif hasattr(self.target.llm, "transformer") and hasattr(self.target.llm.transformer, "h"):
             layer_container = self.target.llm.transformer.h
+            final_norm = getattr(self.target.llm.transformer, "ln_f", None)
         else:
             raise ValueError("Unsupported SALM target structure for DFlash hidden-state capture")
         if isinstance(layer_container, nn.ModuleDict):
@@ -426,8 +497,14 @@ def hook(_module, _args, output):
 
             return hook
 
+        final_layer_id = len(layers) - 1
         for layer_id in self.target_layer_ids:
-            handles.append(layers[layer_id].register_forward_hook(make_hook(layer_id)))
+            if layer_id == final_layer_id:
+                if final_norm is None:
+                    raise ValueError("The final DFlash target-layer tap requires a discoverable target final norm")
+                handles.append(final_norm.register_forward_hook(make_hook(layer_id)))
+            else:
+                handles.append(layers[layer_id].register_forward_hook(make_hook(layer_id)))
 
         forward_kwargs = {
             "inputs_embeds": inputs["input_embeddings"],
@@ -470,24 +547,38 @@ def _run_batch(self, batch: dict[str, torch.Tensor]):
         )
 
     def _globally_normalized_loss(self, metrics) -> torch.Tensor:
-        """Weight a local DFlash mean by the global draft-DP loss denominator.
+        """Weight local draft-loss means by their global draft-DP denominators.
 
         FSDP averages gradients across its process group. Multiplying the local
         weighted-loss numerator by ``dp_size / global_weight`` therefore yields
-        the true global weighted mean even when ranks sample different numbers of
-        valid anchors or supervised block positions.
+        the true global weighted mean even when ranks sample different numbers
+        of valid anchors. DFlash2's backbone and selector terms use different
+        valid sets, so each term must be normalized independently.
         """
-        local_weight = metrics.loss_weight.to(device=metrics.loss.device, dtype=metrics.loss.dtype)
+        loss_terms = self._loss_terms(metrics)
         if not (self._draft_dp_size > 1 and torch.distributed.is_available() and torch.distributed.is_initialized()):
-            return metrics.loss
-
-        global_weight = local_weight.detach().clone()
-        torch.distributed.all_reduce(
-            global_weight,
-            op=torch.distributed.ReduceOp.SUM,
-            group=self._draft_dp_group,
-        )
-        return metrics.loss * local_weight * self._draft_dp_size / global_weight.clamp_min(1.0e-6)
+            return sum(loss for loss, _weight in loss_terms)
+
+        normalized_terms = []
+        for loss, weight in loss_terms:
+            local_weight = weight.to(device=loss.device, dtype=loss.dtype)
+            global_weight = local_weight.detach().clone()
+            torch.distributed.all_reduce(
+                global_weight,
+                op=torch.distributed.ReduceOp.SUM,
+                group=self._draft_dp_group,
+            )
+            normalized_terms.append(loss * local_weight * self._draft_dp_size / global_weight.clamp_min(1.0e-6))
+        return sum(normalized_terms)
+
+    def _loss_terms(self, metrics) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
+        """Return differentiable loss terms paired with their local denominators."""
+        if self.dflash_variant == "dflash2":
+            return (
+                (metrics.base_loss, metrics.loss_weight),
+                (self.selector_loss_weight * metrics.selector_loss, metrics.selector_loss_denominator),
+            )
+        return ((metrics.loss, metrics.loss_weight),)
 
     def training_step(self, batch, batch_idx):
         batches = list(batch.values()) if isinstance(batch, dict) and "input_ids" not in batch else [batch]
@@ -515,6 +606,12 @@ def training_step(self, batch, batch_idx):
             self.log("train/dflash_loss", metrics.loss, on_step=True, prog_bar=True)
             self.log("train/dflash_accuracy", metrics.accuracy, on_step=True)
             self.log("train/accept_len", metrics.accept_len, on_step=True)
+            if self.dflash_variant == "dflash2":
+                self.log("train/dflash_base_loss", metrics.base_loss, on_step=True)
+                self.log("train/dflash_selector_loss", metrics.selector_loss, on_step=True)
+                self.log("train/dflash_base_accuracy", metrics.base_accuracy, on_step=True)
+                self.log("train/dflash_base_accept_len", metrics.base_accept_len, on_step=True)
+                self.log("train/dflash_candidate_recall", metrics.candidate_recall, on_step=True)
         if not losses:
             # Every rank takes the same synchronized skip branches above. Lightning
             # rejects ``None`` from ``training_step`` under distributed automatic
@@ -554,18 +651,32 @@ def validation_step(self, batch, batch_idx) -> None:
             # float64 so a single stacked all-reduce remains exact for counts.
             metric_dtype = torch.float64
             metric_device = metrics.loss.device
-            self._partial_val_metrics[dataset_name].append(
-                torch.stack(
-                    [
-                        metrics.loss.detach().to(dtype=metric_dtype) * metrics.loss_weight.to(metric_dtype),
-                        metrics.loss_weight.to(dtype=metric_dtype, device=metric_device),
-                        metrics.correct_tokens.to(dtype=metric_dtype, device=metric_device),
-                        metrics.valid_tokens.to(dtype=metric_dtype, device=metric_device),
-                        metrics.accept_len_sum.to(dtype=metric_dtype, device=metric_device),
-                        metrics.valid_blocks.to(dtype=metric_dtype, device=metric_device),
-                    ]
-                )
-            )
+            if self.dflash_variant == "dflash2":
+                selector_weight = metrics.selector_loss_denominator.to(dtype=metric_dtype, device=metric_device)
+                valid_tokens = metrics.valid_tokens.to(dtype=metric_dtype, device=metric_device)
+                metric_sums = [
+                    metrics.base_loss.detach().to(dtype=metric_dtype) * metrics.loss_weight.to(metric_dtype),
+                    metrics.loss_weight.to(dtype=metric_dtype, device=metric_device),
+                    metrics.selector_loss.detach().to(dtype=metric_dtype) * selector_weight,
+                    selector_weight,
+                    metrics.correct_tokens.to(dtype=metric_dtype, device=metric_device),
+                    valid_tokens,
+                    metrics.accept_len_sum.to(dtype=metric_dtype, device=metric_device),
+                    metrics.valid_blocks.to(dtype=metric_dtype, device=metric_device),
+                    metrics.base_correct_tokens.to(dtype=metric_dtype, device=metric_device),
+                    metrics.base_accept_len_sum.to(dtype=metric_dtype, device=metric_device),
+                    metrics.candidate_recall.to(dtype=metric_dtype, device=metric_device) * valid_tokens,
+                ]
+            else:
+                metric_sums = [
+                    metrics.loss.detach().to(dtype=metric_dtype) * metrics.loss_weight.to(metric_dtype),
+                    metrics.loss_weight.to(dtype=metric_dtype, device=metric_device),
+                    metrics.correct_tokens.to(dtype=metric_dtype, device=metric_device),
+                    metrics.valid_tokens.to(dtype=metric_dtype, device=metric_device),
+                    metrics.accept_len_sum.to(dtype=metric_dtype, device=metric_device),
+                    metrics.valid_blocks.to(dtype=metric_dtype, device=metric_device),
+                ]
+            self._partial_val_metrics[dataset_name].append(torch.stack(metric_sums))
 
     def on_validation_epoch_end(self) -> None:
         all_sums = []
@@ -582,11 +693,39 @@ def on_validation_epoch_end(self) -> None:
         self._partial_val_metrics.clear()
 
     def _log_validation_metrics(self, metric_sums: torch.Tensor, suffix: str = "") -> None:
-        loss_sum, loss_weight, correct, valid, accept_sum, valid_blocks = metric_sums
+        if self.dflash_variant == "dflash2":
+            (
+                base_loss_sum,
+                base_loss_weight,
+                selector_loss_sum,
+                selector_loss_weight,
+                correct,
+                valid,
+                accept_sum,
+                valid_blocks,
+                base_correct,
+                base_accept_sum,
+                candidate_hits,
+            ) = metric_sums
+            base_loss = base_loss_sum / base_loss_weight.clamp_min(1)
+            selector_loss = selector_loss_sum / selector_loss_weight.clamp_min(1)
+            loss = base_loss + self.selector_loss_weight * selector_loss
+            self.log(f"val/dflash_base_loss{suffix}", base_loss, on_epoch=True)
+            self.log(f"val/dflash_selector_loss{suffix}", selector_loss, on_epoch=True)
+            self.log(f"val/dflash_base_accuracy{suffix}", base_correct / valid.clamp_min(1), on_epoch=True)
+            self.log(
+                f"val/dflash_base_accept_len{suffix}",
+                base_accept_sum / valid_blocks.clamp_min(1),
+                on_epoch=True,
+            )
+            self.log(f"val/dflash_candidate_recall{suffix}", candidate_hits / valid.clamp_min(1), on_epoch=True)
+        else:
+            loss_sum, loss_weight, correct, valid, accept_sum, valid_blocks = metric_sums
+            loss = loss_sum / loss_weight.clamp_min(1)
         accuracy = correct / valid.clamp_min(1)
         self.log(
             f"val/dflash_loss{suffix}",
-            loss_sum / loss_weight.clamp_min(1),
+            loss,
             on_epoch=True,
         )
         self.log(f"val/dflash_accuracy{suffix}", accuracy, on_epoch=True)
diff --git a/pyproject.toml b/pyproject.toml
index e48b9c6657db..0082213d2677 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -472,10 +472,10 @@ deep-ep = ["setuptools"]
 # --- uv configuration ---
 # Keep Torch wheel indexes explicit per CUDA extra. Automodel's upstream source
 # metadata is mirrored below so this repository controls the CUDA wheel index.
-# Draft-only stacked pin to Automodel PR #3673; replace with its upstream merge
-# commit before marking the dependent NeMo PR ready for review.
+# Draft-only immutable pin to the implementation reviewed in Automodel PR #3673;
+# replace with its upstream merge commit before marking the dependent NeMo PR ready.
 [tool.uv.sources]
-nemo_automodel = { git = "https://github.com/NVIDIA-NeMo/Automodel.git", rev = "refs/pull/3673/head" }
+nemo_automodel = { git = "https://github.com/Slyne/Automodel.git", rev = "426711b8336d563c37c7ec4a9b97b1bb6dd86aca" }
 megatron-fsdp = { git = "https://github.com/yuhezhang-ai/Megatron-LM.git", rev = "455389c480af6b3acdca74c7830c68b3274eb083", subdirectory = "megatron/core/distributed/fsdp/src" }
 deep_ep = { git = "https://github.com/deepseek-ai/DeepEP.git", rev = "7febc6e25660af0f54d95dd781ecdcd62265ecca" }
 torch = [
diff --git a/tests/collections/speechlm2/test_salm_dflash.py b/tests/collections/speechlm2/test_salm_dflash.py
index ff989ec6572c..92fbb4f45cd0 100644
--- a/tests/collections/speechlm2/test_salm_dflash.py
+++ b/tests/collections/speechlm2/test_salm_dflash.py
@@ -26,6 +26,10 @@
 pytestmark = pytest.mark.unit
 
 from nemo_automodel.components.loss.dllm_loss import DFlashDecayLoss  # noqa: E402
+from nemo_automodel.components.speculative.dflash.draft_qwen3 import Qwen3DFlashDraftModel  # noqa: E402
+from nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2 import (  # noqa: E402
+    Qwen3DFlash2DraftModel,
+)
 
 from nemo.collections.speechlm2.parts import dflash as salm_dflash  # noqa: E402
 
@@ -264,9 +268,172 @@ def test_build_draft_config_applies_explicit_architecture_and_layer_taps():
     assert draft_config.num_hidden_layers == 2
     assert draft_config.intermediate_size == 96
     assert draft_config.dflash_config == {
+        "block_size": 8,
+        "mask_token_id": 18,
+        "target_layer_ids": [1, 6],
+    }
+    assert draft_config.architectures == ["Qwen3DFlashDraftModel"]
+    assert draft_config.is_causal is False
+
+
+def test_build_dflash2_config_matches_automodel_recipe_fields():
+    target_config = Qwen3Config(
+        hidden_size=64,
+        intermediate_size=128,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        num_hidden_layers=8,
+        head_dim=16,
+        vocab_size=128,
+    )
+
+    draft_config, layer_ids = salm_dflash._build_draft_config(
+        target_config,
+        {
+            "variant": "dflash2",
+            "draft_num_hidden_layers": 2,
+            "target_layer_ids": [1, 6],
+            "conv_kernel_size": 2,
+            "conv_group_size": 16,
+            "selector_rank": 32,
+            "selector_top_k": 8,
+            "draft_sliding_window": 32,
+        },
+        block_size=8,
+        mask_token_id=18,
+    )
+
+    assert layer_ids == [1, 6]
+    assert draft_config.architectures == ["Qwen3DFlash2DraftModel"]
+    assert draft_config.layer_types == ["sliding_attention", "sliding_attention"]
+    assert draft_config.sliding_window == 32
+    assert draft_config.use_sliding_window is True
+    assert draft_config.dflash_config == {
+        "block_size": 8,
         "mask_token_id": 18,
         "target_layer_ids": [1, 6],
+        "conv_kernel_size": 2,
+        "conv_group_size": 16,
+        "selector_rank": 32,
+        "selector_top_k": 8,
     }
+    draft = Qwen3DFlash2DraftModel(draft_config)
+    assert draft.candidate_selector.top_k == 8
+
+
+def test_dflash2_rejects_fused_linear_ce():
+    with pytest.raises(ValueError, match="use_fused_linear_ce"):
+        salm_dflash.SALMDFlashModule(
+            nn.Linear(1, 1),
+            {"dflash": {"variant": "dflash2", "mask_token_id": 18, "use_fused_linear_ce": True}},
+        )
+
+
+class _TrainerTargetLLM(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.config = Qwen3Config(
+            hidden_size=32,
+            intermediate_size=64,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            num_hidden_layers=6,
+            head_dim=8,
+            vocab_size=64,
+        )
+        self.embed_tokens = nn.Embedding(64, 32)
+        self.lm_head = nn.Linear(32, 64, bias=False)
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+
+class _TrainerTarget(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.llm = _TrainerTargetLLM()
+
+
+@pytest.mark.parametrize(
+    "variant,trainer_name",
+    [("dflash", "DFlashTrainerModule"), ("dflash2", "DFlash2TrainerModule")],
+)
+def test_create_trainer_module_selects_configured_variant(monkeypatch, variant, trainer_name):
+    base_factory = Mock(return_value=object())
+    dflash2_factory = Mock(return_value=object())
+    monkeypatch.setattr(salm_dflash, "DFlashTrainerModule", base_factory)
+    monkeypatch.setattr(salm_dflash, "DFlash2TrainerModule", dflash2_factory)
+    module = salm_dflash.SALMDFlashModule(
+        _TrainerTarget(),
+        {
+            "dflash": {
+                "variant": variant,
+                "mask_token_id": 18,
+                "max_total_anchors": 64,
+                "selector_loss_weight": 0.25,
+                "use_fused_linear_ce": False,
+            }
+        },
+    )
+    module.draft_model = nn.Linear(32, 32)
+
+    result = module._create_trainer_module()
+
+    selected = dflash2_factory if trainer_name == "DFlash2TrainerModule" else base_factory
+    unselected = base_factory if trainer_name == "DFlash2TrainerModule" else dflash2_factory
+    assert result is selected.return_value
+    expected_draft_cls = Qwen3DFlash2DraftModel if variant == "dflash2" else Qwen3DFlashDraftModel
+    assert module._draft_model_class() is expected_draft_cls
+    unselected.assert_not_called()
+    kwargs = selected.call_args.kwargs
+    assert kwargs["draft_model"] is module.draft_model
+    assert kwargs["max_total_anchors"] == 64
+    if variant == "dflash2":
+        assert kwargs["selector_loss_weight"] == pytest.approx(0.25)
+        assert "use_fused_linear_ce" not in kwargs
+    else:
+        assert kwargs["use_fused_linear_ce"] is False
+
+
+def test_dflash2_salm_components_run_forward_and_train_selector():
+    torch.manual_seed(7)
+    target = _TrainerTarget()
+    dflash_config = {
+        "variant": "dflash2",
+        "mask_token_id": 63,
+        "block_size": 4,
+        "draft_num_hidden_layers": 2,
+        "target_layer_ids": [1, 4],
+        "conv_group_size": 8,
+        "selector_rank": 16,
+        "selector_top_k": 64,
+        "num_anchors": 2,
+        "max_total_anchors": 2,
+        "attention_backend": "sdpa",
+        "activation_checkpointing": False,
+    }
+    module = salm_dflash.SALMDFlashModule(target, {"dflash": dflash_config})
+    draft_config, module.target_layer_ids = salm_dflash._build_draft_config(
+        target.llm.config,
+        dflash_config,
+        block_size=4,
+        mask_token_id=63,
+    )
+    draft_config._attn_implementation = "sdpa"
+    module.draft_model = module._draft_model_class()(draft_config)
+    module.trainer_module = module._create_trainer_module()
+    input_ids = torch.randint(0, 63, (1, 8))
+    hidden_states = torch.randn(1, 8, 64)
+
+    metrics = module.trainer_module(input_ids=input_ids, hidden_states=hidden_states, loss_mask=torch.ones(1, 8))
+    metrics.loss.backward()
+
+    assert torch.isfinite(metrics.loss)
+    assert metrics.selector_loss.item() > 0
+    assert module.draft_model.candidate_selector.successor_codebook.grad.abs().sum() > 0
 
 
 def test_build_draft_config_rejects_managed_overrides():
@@ -281,7 +448,7 @@ def test_build_draft_config_rejects_managed_overrides():
         )
 
 
-def test_salm_automodel_dflash_defaults_match_nemotron_3_5_lightning():
+def test_salm_automodel_dflash2_defaults_match_nemotron_3_5_lightning():
     cfg = OmegaConf.load(REPO_ROOT / "examples/speechlm2/conf/salm_automodel.yaml")
     dflash_cfg = OmegaConf.to_container(cfg.dflash, resolve=True)
     target_config = Qwen3Config(
@@ -302,13 +469,14 @@ def test_salm_automodel_dflash_defaults_match_nemotron_3_5_lightning():
     )
 
     assert dflash_cfg["enabled"] is False
+    assert dflash_cfg["variant"] == "dflash2"
     assert dflash_cfg["block_size"] == 8
     assert dflash_cfg["num_anchors"] == 512
     assert dflash_cfg["max_total_anchors"] == 512
     assert dflash_cfg["loss_decay_gamma"] == pytest.approx(4.0)
     assert dflash_cfg["attention_backend"] == "flex_attention"
     assert dflash_cfg["activation_checkpointing"] is True
-    assert dflash_cfg["use_fused_linear_ce"] is True
+    assert dflash_cfg["use_fused_linear_ce"] is False
     assert dflash_cfg["linear_ce_chunk_size"] == 256
     assert draft_config.num_hidden_layers == 6
     assert draft_config.hidden_size == 2688
@@ -326,10 +494,16 @@ def test_salm_automodel_dflash_defaults_match_nemotron_3_5_lightning():
     }
     assert target_layer_ids == [1, 5, 19, 29, 41, 51]
     assert draft_config.dflash_config == {
+        "block_size": 8,
         "mask_token_id": 990,
         "target_layer_ids": [1, 5, 19, 29, 41, 51],
+        "conv_kernel_size": 2,
+        "conv_group_size": 16,
+        "selector_rank": 256,
+        "selector_top_k": 16,
     }
     assert draft_config.block_size == 8
+    assert draft_config.architectures == ["Qwen3DFlash2DraftModel"]
 
 
 class _TargetLLM(nn.Module):
@@ -337,6 +511,7 @@ def __init__(self):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(1))
         self.layers = nn.ModuleList([nn.Identity(), nn.Identity(), nn.Identity()])
+        self.norm = _AddConstant(10.0)
         self.calls = []
 
     def forward(
@@ -361,6 +536,7 @@ def forward(
         hidden = inputs_embeds
         for index, layer in enumerate(self.layers, start=1):
             hidden = layer(hidden + index)
+        hidden = self.norm(hidden)
         return SimpleNamespace(hidden_states=(hidden,))
 
 
@@ -370,6 +546,15 @@ def __init__(self):
         self.llm = _TargetLLM()
 
 
+class _AddConstant(nn.Module):
+    def __init__(self, value: float):
+        super().__init__()
+        self.value = value
+
+    def forward(self, inputs):
+        return inputs + self.value
+
+
 class _MinimalTargetLLM(nn.Module):
     """Target whose explicit forward rejects every optional HF-style kwarg."""
 
@@ -377,6 +562,7 @@ def __init__(self):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(1))
         self.layers = nn.ModuleList([nn.Identity(), nn.Identity()])
+        self.norm = nn.Identity()
         self.calls = []
 
     def forward(self, *, inputs_embeds, attention_mask):
@@ -384,7 +570,7 @@ def forward(self, *, inputs_embeds, attention_mask):
         hidden = inputs_embeds
         for index, layer in enumerate(self.layers, start=1):
             hidden = layer(hidden + index)
-        return hidden
+        return self.norm(hidden)
 
 
 def test_target_hidden_states_uses_audio_embeddings_and_skips_logits():
@@ -399,7 +585,7 @@ def test_target_hidden_states_uses_audio_embeddings_and_skips_logits():
 
     assert hidden.shape == (2, 5, 8)
     assert torch.allclose(hidden[..., :4], inputs["input_embeddings"] + 1)
-    assert torch.allclose(hidden[..., 4:], inputs["input_embeddings"] + 6)
+    assert torch.allclose(hidden[..., 4:], inputs["input_embeddings"] + 16)
     assert module.target.llm.calls == [
         {
             "attention_mask": inputs["attention_mask"],
@@ -477,6 +663,41 @@ def fake_all_reduce(value, *, op, group):
     assert local_loss.grad.item() == pytest.approx(0.6)
 
 
+def test_dflash2_globally_normalizes_base_and_selector_terms_separately(monkeypatch):
+    module = salm_dflash.SALMDFlashModule(
+        nn.Linear(1, 1),
+        {"dflash": {"variant": "dflash2", "mask_token_id": 18, "selector_loss_weight": 0.5}},
+    )
+    module._draft_dp_size = 2
+    module._draft_dp_group = object()
+    monkeypatch.setattr(salm_dflash.torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(salm_dflash.torch.distributed, "is_initialized", lambda: True)
+    global_weights = iter((10.0, 8.0))
+
+    def fake_all_reduce(value, *, op, group):
+        assert op == salm_dflash.torch.distributed.ReduceOp.SUM
+        assert group is module._draft_dp_group
+        value.fill_(next(global_weights))
+
+    monkeypatch.setattr(salm_dflash.torch.distributed, "all_reduce", fake_all_reduce)
+    base_loss = torch.tensor(2.0, requires_grad=True)
+    selector_loss = torch.tensor(4.0, requires_grad=True)
+    metrics = SimpleNamespace(
+        loss=base_loss + 0.5 * selector_loss,
+        loss_weight=torch.tensor(3.0),
+        base_loss=base_loss,
+        selector_loss=selector_loss,
+        selector_loss_denominator=torch.tensor(2.0),
+    )
+
+    loss = module._globally_normalized_loss(metrics)
+    loss.backward()
+
+    assert loss.item() == pytest.approx(2.2)
+    assert base_loss.grad.item() == pytest.approx(0.6)
+    assert selector_loss.grad.item() == pytest.approx(0.25)
+
+
 def test_dflash_loss_times_weight_recovers_decay_weighted_numerator():
     torch.manual_seed(7)
     block_size = 4
@@ -590,6 +811,24 @@ def test_aggregate_validation_accuracy_preserves_default_checkpoint_monitor(
     log.assert_any_call("val_acc", torch.tensor(0.5), on_epoch=True)
 
 
+def test_dflash2_validation_uses_separate_loss_denominators_and_selector_metrics(monkeypatch):
+    module = salm_dflash.SALMDFlashModule(
+        nn.Linear(1, 1),
+        {"dflash": {"variant": "dflash2", "mask_token_id": 18, "selector_loss_weight": 0.5}},
+    )
+    log = Mock()
+    monkeypatch.setattr(module, "log", log)
+
+    module._log_validation_metrics(torch.tensor([8.0, 4.0, 6.0, 3.0, 3.0, 6.0, 5.0, 2.0, 2.0, 4.0, 5.0]))
+
+    log.assert_any_call("val/dflash_loss", torch.tensor(3.0), on_epoch=True)
+    log.assert_any_call("val/dflash_selector_loss", torch.tensor(2.0), on_epoch=True)
+    log.assert_any_call("val/dflash_accuracy", torch.tensor(0.5), on_epoch=True)
+    log.assert_any_call("val/dflash_base_accept_len", torch.tensor(2.0), on_epoch=True)
+    log.assert_any_call("val/dflash_candidate_recall", torch.tensor(5.0 / 6.0), on_epoch=True)
+    log.assert_any_call("val_acc", torch.tensor(0.5), on_epoch=True)
+
+
 def test_state_dict_hook_keeps_only_draft_parameters():
     module = SimpleNamespace(_CHECKPOINT_STATE_PREFIX="draft_model.")
     state_dict = {
diff --git a/uv.lock b/uv.lock
index b2d26a681d4a..2942fc2db698 100644
--- a/uv.lock
+++ b/uv.lock
@@ -4180,7 +4180,7 @@ wheels = [
 [[package]]
 name = "nemo-automodel"
 version = "0.5.0+426711b8"
-source = { git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=refs%2Fpull%2F3673%2Fhead#426711b8336d563c37c7ec4a9b97b1bb6dd86aca" }
+source = { git = "https://github.com/Slyne/Automodel.git?rev=426711b8336d563c37c7ec4a9b97b1bb6dd86aca#426711b8336d563c37c7ec4a9b97b1bb6dd86aca" }
 dependencies = [
     { name = "datasets" },
     { name = "flashoptim" },
@@ -4633,9 +4633,9 @@ requires-dist = [
     { name = "matplotlib", marker = "extra == 'audio'" },
     { name = "matplotlib", marker = "extra == 'speechlm2'" },
     { name = "matplotlib", marker = "extra == 'tts'" },
-    { name = "nemo-automodel", marker = "extra == 'all'", git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=refs%2Fpull%2F3673%2Fhead" },
-    { name = "nemo-automodel", marker = "extra == 'speechlm2'", git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=refs%2Fpull%2F3673%2Fhead" },
-    { name = "nemo-automodel", marker = "extra == 'speechlm2-only'", git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=refs%2Fpull%2F3673%2Fhead" },
+    { name = "nemo-automodel", marker = "extra == 'all'", git = "https://github.com/Slyne/Automodel.git?rev=426711b8336d563c37c7ec4a9b97b1bb6dd86aca" },
+    { name = "nemo-automodel", marker = "extra == 'speechlm2'", git = "https://github.com/Slyne/Automodel.git?rev=426711b8336d563c37c7ec4a9b97b1bb6dd86aca" },
+    { name = "nemo-automodel", marker = "extra == 'speechlm2-only'", git = "https://github.com/Slyne/Automodel.git?rev=426711b8336d563c37c7ec4a9b97b1bb6dd86aca" },
     { name = "nemo-text-processing", marker = "'aarch' not in platform_machine and 'arm' not in platform_machine and sys_platform != 'darwin' and extra == 'all'" },
     { name = "nemo-text-processing", marker = "'aarch' not in platform_machine and 'arm' not in platform_machine and sys_platform != 'darwin' and extra == 'speechlm2'" },
     { name = "nemo-text-processing", marker = "'aarch' not in platform_machine and 'arm' not in platform_machine and sys_platform != 'darwin' and extra == 'tts'" },

From 0cf301f8b65fd11a0ee4535bfa9a532bf9fe6a70 Mon Sep 17 00:00:00 2001
From: SlyneD <slyned@nvidia.com>
Date: Wed, 26 Aug 2026 14:29:24 -0700
Subject: [PATCH 7/8] fix(speechlm2): capture DFlash targets before final norm

Signed-off-by: SlyneD <slyned@nvidia.com>
---
 nemo/collections/speechlm2/parts/dflash.py    | 24 ++++++++++---------
 .../collections/speechlm2/test_salm_dflash.py |  7 ++++--
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/nemo/collections/speechlm2/parts/dflash.py b/nemo/collections/speechlm2/parts/dflash.py
index 8bd02ec8977a..4e67518b7d90 100644
--- a/nemo/collections/speechlm2/parts/dflash.py
+++ b/nemo/collections/speechlm2/parts/dflash.py
@@ -471,16 +471,24 @@ def _prepare_batch(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tens
 
     @torch.no_grad()
     def _target_hidden_states(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor:
-        """Run the frozen audio-conditioned target and concatenate configured layers."""
+        """Run the frozen target and concatenate pre-final-norm decoder-block outputs.
+
+        Args:
+            inputs: Mapping containing ``input_embeddings`` with shape
+                ``[batch, sequence, hidden]`` and ``attention_mask`` with shape
+                ``[batch, sequence]``.
+
+        Returns:
+            Tensor of shape ``[batch, sequence, selected_layers * hidden]``.
+            Each feature is the raw output of the configured decoder block,
+            before any separate model-level final normalization.
+        """
         if hasattr(self.target.llm, "model") and hasattr(self.target.llm.model, "layers"):
             layer_container = self.target.llm.model.layers
-            final_norm = getattr(self.target.llm.model, "norm", None)
         elif hasattr(self.target.llm, "layers"):
             layer_container = self.target.llm.layers
-            final_norm = getattr(self.target.llm, "norm", None)
         elif hasattr(self.target.llm, "transformer") and hasattr(self.target.llm.transformer, "h"):
             layer_container = self.target.llm.transformer.h
-            final_norm = getattr(self.target.llm.transformer, "ln_f", None)
         else:
             raise ValueError("Unsupported SALM target structure for DFlash hidden-state capture")
         if isinstance(layer_container, nn.ModuleDict):
@@ -497,14 +505,8 @@ def hook(_module, _args, output):
 
             return hook
 
-        final_layer_id = len(layers) - 1
         for layer_id in self.target_layer_ids:
-            if layer_id == final_layer_id:
-                if final_norm is None:
-                    raise ValueError("The final DFlash target-layer tap requires a discoverable target final norm")
-                handles.append(final_norm.register_forward_hook(make_hook(layer_id)))
-            else:
-                handles.append(layers[layer_id].register_forward_hook(make_hook(layer_id)))
+            handles.append(layers[layer_id].register_forward_hook(make_hook(layer_id)))
 
         forward_kwargs = {
             "inputs_embeds": inputs["input_embeddings"],
diff --git a/tests/collections/speechlm2/test_salm_dflash.py b/tests/collections/speechlm2/test_salm_dflash.py
index 92fbb4f45cd0..e2ed7bcfbd37 100644
--- a/tests/collections/speechlm2/test_salm_dflash.py
+++ b/tests/collections/speechlm2/test_salm_dflash.py
@@ -573,7 +573,7 @@ def forward(self, *, inputs_embeds, attention_mask):
         return self.norm(hidden)
 
 
-def test_target_hidden_states_uses_audio_embeddings_and_skips_logits():
+def test_target_hidden_states_uses_pre_final_norm_block_outputs_and_skips_logits():
     module = salm_dflash.SALMDFlashModule(_TargetModel(), {"dflash": {"mask_token_id": 18}})
     module.target_layer_ids = [0, 2]
     inputs = {
@@ -585,7 +585,10 @@ def test_target_hidden_states_uses_audio_embeddings_and_skips_logits():
 
     assert hidden.shape == (2, 5, 8)
     assert torch.allclose(hidden[..., :4], inputs["input_embeddings"] + 1)
-    assert torch.allclose(hidden[..., 4:], inputs["input_embeddings"] + 16)
+    # Block 2 contributes +3 after blocks 0 and 1 contributed +1 and +2.
+    # The separate final norm contributes +10 to the model output, but it must
+    # not alter DFlash's captured decoder-block feature.
+    assert torch.allclose(hidden[..., 4:], inputs["input_embeddings"] + 6)
     assert module.target.llm.calls == [
         {
             "attention_mask": inputs["attention_mask"],

From 25f3854037b7c6e6d8319d9539b4f1b4e44bf675 Mon Sep 17 00:00:00 2001
From: SlyneD <slyned@nvidia.com>
Date: Fri, 4 Sep 2026 09:52:37 -0700
Subject: [PATCH 8/8] feat(speechlm2): support packed SALM DFlash training

Signed-off-by: SlyneD <slyned@nvidia.com>
---
 nemo/collections/speechlm2/parts/dflash.py    | 172 +++++++-
 .../speechlm2/parts/packed_sequences.py       | 225 ++++++++++
 .../collections/speechlm2/test_salm_dflash.py | 405 +++++++++++++++++-
 3 files changed, 769 insertions(+), 33 deletions(-)

diff --git a/nemo/collections/speechlm2/parts/dflash.py b/nemo/collections/speechlm2/parts/dflash.py
index 4e67518b7d90..1d4a78ddbacd 100644
--- a/nemo/collections/speechlm2/parts/dflash.py
+++ b/nemo/collections/speechlm2/parts/dflash.py
@@ -19,6 +19,7 @@
 import inspect
 from collections import defaultdict
 from collections.abc import Sequence
+from pathlib import Path
 
 import torch
 from lightning import LightningModule
@@ -44,6 +45,10 @@
     encode_audio_with_cp_distribution,
     get_perception_fsdp_group,
 )
+from nemo.collections.speechlm2.parts.packed_sequences import (
+    _validate_packed_dflash_inputs,
+    pack_audio_for_dflash,
+)
 from nemo.core.classes.common import safe_instantiate
 
 _DRAFT_CONFIG_MANAGED_KEYS = {
@@ -107,16 +112,29 @@ def _synchronize_ep_group_before_target_forward(moe_mesh) -> None:
         torch.distributed.barrier(group=ep_mesh.get_group())
 
 
-def _has_valid_dflash_anchors(loss_mask: torch.Tensor, block_size: int) -> bool:
-    """Mirror Automodel's unpacked DFlash anchor-validity predicate.
+def _has_valid_dflash_anchors(
+    loss_mask: torch.Tensor,
+    block_size: int,
+    doc_remaining: torch.Tensor | None = None,
+) -> bool:
+    """Mirror Automodel's DFlash anchor-validity predicate.
 
-    SALM DFlash rejects packed sequences, so ``DFlashTrainerModule`` considers
-    an anchor valid exactly when its own position is supervised and it lies no
-    later than ``seq_len - block_size``. Following block positions may be masked;
-    they affect the loss denominator but not anchor validity.
+    Under packing, a valid anchor must also leave ``block_size - 1`` real
+    tokens in the same document. Keeping this precheck identical to Automodel
+    makes the subsequent skip decision rank-synchronous.
     """
+    if loss_mask.ndim != 2:
+        raise ValueError(f"DFlash loss_mask must have shape [B, S], got {tuple(loss_mask.shape)}")
+    if doc_remaining is not None and doc_remaining.shape != loss_mask.shape:
+        raise ValueError(
+            "Packed DFlash doc_remaining must match loss_mask shape; "
+            f"got {tuple(doc_remaining.shape)} and {tuple(loss_mask.shape)}"
+        )
     max_anchor = max(loss_mask.shape[1] - block_size, 0)
-    return bool((loss_mask[:, : max_anchor + 1] > 0.5).any().item())
+    valid = loss_mask[:, : max_anchor + 1] > 0.5
+    if doc_remaining is not None:
+        valid = valid & (doc_remaining[:, : max_anchor + 1] >= block_size - 1)
+    return bool(valid.any().item())
 
 
 def _validate_dflash_parallelism(mesh_context) -> None:
@@ -435,13 +453,22 @@ def _audio_embeddings(self, batch: dict[str, torch.Tensor]) -> list[torch.Tensor
         )
 
     def _prepare_batch(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
-        if self.target.cfg.get("packed_sequences", False):
-            raise NotImplementedError("SALM DFlash currently requires model.packed_sequences=false")
-
+        """Expand audio and prepare an unshifted BSHD or packed THD DFlash stream."""
         input_ids = batch["input_ids"]
         audio_embeddings = self._audio_embeddings(batch)
         text_ids = torch.where(input_ids == self.target.audio_locator_tag_id, 0, input_ids)
         text_embeddings = self.target._embed_tokens(text_ids)
+        if self.target.cfg.get("packed_sequences", False):
+            return pack_audio_for_dflash(
+                input_ids=input_ids,
+                embeds=text_embeddings,
+                loss_mask=batch["loss_mask"],
+                replacements=audio_embeddings,
+                padding_id=self.target.text_pad_id,
+                placeholder_id=self.target.audio_locator_tag_id,
+                mask_token_id=self.mask_token_id,
+            )
+
         target_ids = input_ids.where(batch["loss_mask"], -100)
         input_embeddings, target_ids, attention_mask = replace_placeholders_and_build_targets(
             input_ids=input_ids,
@@ -474,15 +501,21 @@ def _target_hidden_states(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor
         """Run the frozen target and concatenate pre-final-norm decoder-block outputs.
 
         Args:
-            inputs: Mapping containing ``input_embeddings`` with shape
-                ``[batch, sequence, hidden]`` and ``attention_mask`` with shape
-                ``[batch, sequence]``.
+            inputs: Mapping containing BSHD ``input_embeddings`` with shape
+                ``[batch, sequence, hidden]`` and ``attention_mask``, or packed
+                THD embeddings ``[tokens, hidden]`` plus ``qkv_format``,
+                ``cu_seqlens``, ``position_ids``, and ``max_seqlen``.
 
         Returns:
             Tensor of shape ``[batch, sequence, selected_layers * hidden]``.
+            THD hook outputs are normalized to one synthetic packed batch row.
             Each feature is the raw output of the configured decoder block,
             before any separate model-level final normalization.
         """
+        packing_fields = ("qkv_format", "cu_seqlens", "position_ids", "seq_lens", "doc_remaining", "max_seqlen")
+        is_packed = any(name in inputs for name in packing_fields)
+        if is_packed:
+            _validate_packed_dflash_inputs(inputs)
         if hasattr(self.target.llm, "model") and hasattr(self.target.llm.model, "layers"):
             layer_container = self.target.llm.model.layers
         elif hasattr(self.target.llm, "layers"):
@@ -516,12 +549,22 @@ def hook(_module, _args, output):
         accepts_extra_kwargs = any(
             parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in forward_parameters.values()
         )
-        for name, value in {
+        optional_kwargs = {
             "output_hidden_states": False,
             "use_cache": False,
             "return_dict": True,
             "compute_logits": False,
-        }.items():
+        }
+        if is_packed:
+            optional_kwargs.update(
+                {
+                    "qkv_format": "thd",
+                    "cu_seqlens": inputs["cu_seqlens"],
+                    "position_ids": inputs["position_ids"],
+                    "max_seqlen": inputs["max_seqlen"],
+                }
+            )
+        for name, value in optional_kwargs.items():
             if accepts_extra_kwargs or name in forward_parameters:
                 forward_kwargs[name] = value
         try:
@@ -531,22 +574,59 @@ def hook(_module, _args, output):
                 handle.remove()
         if len(captured) != len(self.target_layer_ids):
             raise RuntimeError(f"Expected {len(self.target_layer_ids)} captured target layers, got {sorted(captured)}")
-        return torch.cat([captured[layer_id] for layer_id in self.target_layer_ids], dim=-1)
+        hidden_states = torch.cat([captured[layer_id] for layer_id in self.target_layer_ids], dim=-1)
+        if is_packed:
+            if hidden_states.ndim == 2:
+                hidden_states = hidden_states.unsqueeze(0)
+            if hidden_states.ndim != 3 or hidden_states.shape[0] != 1:
+                raise RuntimeError(
+                    "Packed DFlash target hooks must return [T, H] or [1, T, H], got " f"{tuple(hidden_states.shape)}"
+                )
+            expected_tokens = inputs["input_ids"].shape[1]
+            if hidden_states.shape[1] != expected_tokens:
+                raise RuntimeError(
+                    "Packed DFlash target features are not token aligned: "
+                    f"got {hidden_states.shape[1]} features for {expected_tokens} input IDs"
+                )
+        return hidden_states
 
     def _run_batch(self, batch: dict[str, torch.Tensor]):
         inputs = self._prepare_batch(batch)
+        is_packed = inputs.get("qkv_format") == "thd"
+        # Keep the expanded, non-padding target-input count next to the result
+        # without changing Automodel's public metric dataclasses. The training
+        # loop consumes it immediately after _run_batch returns. For THD this is
+        # the sum of document lengths; for padded BSHD the attention mask is the
+        # equivalent count after audio-frame expansion.
+        self._last_input_token_count = (
+            inputs["seq_lens"].sum() if is_packed else inputs["attention_mask"].sum()
+        ).detach()
+        doc_remaining = inputs["doc_remaining"] if is_packed else None
         if not _all_ranks_agree(
-            _has_valid_dflash_anchors(inputs["loss_mask"], self.block_size),
+            _has_valid_dflash_anchors(
+                inputs["loss_mask"],
+                self.block_size,
+                doc_remaining=doc_remaining,
+            ),
             inputs["loss_mask"].device,
         ):
             raise NoValidAnchorsError("At least one rank has no valid DFlash anchors")
         _synchronize_ep_group_before_target_forward(getattr(self.trainer.strategy, "moe_mesh", None))
         hidden_states = self._target_hidden_states(inputs)
-        return self.trainer_module(
-            input_ids=inputs["input_ids"],
-            hidden_states=hidden_states,
-            loss_mask=inputs["loss_mask"],
-        )
+        trainer_kwargs = {
+            "input_ids": inputs["input_ids"],
+            "hidden_states": hidden_states,
+            "loss_mask": inputs["loss_mask"],
+        }
+        if is_packed:
+            trainer_kwargs.update(
+                {
+                    "position_ids": inputs["position_ids"],
+                    "seq_lens": inputs["seq_lens"],
+                    "doc_remaining": inputs["doc_remaining"],
+                }
+            )
+        return self.trainer_module(**trainer_kwargs)
 
     def _globally_normalized_loss(self, metrics) -> torch.Tensor:
         """Weight local draft-loss means by their global draft-DP denominators.
@@ -585,6 +665,10 @@ def _loss_terms(self, metrics) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
     def training_step(self, batch, batch_idx):
         batches = list(batch.values()) if isinstance(batch, dict) and "input_ids" not in batch else [batch]
         losses = []
+        input_token_counts = []
+        valid_token_counts = []
+        valid_block_counts = []
+        loss_weights = []
         skip_counts = defaultdict(int)
         num_batches = _max_rank_value(len(batches), self.device)
         for dataset_index in range(num_batches):
@@ -605,6 +689,10 @@ def training_step(self, batch, batch_idx):
                 skip_counts["no_valid_anchors"] += 1
                 continue
             losses.append(self._globally_normalized_loss(metrics))
+            input_token_counts.append(self._last_input_token_count)
+            valid_token_counts.append(metrics.valid_tokens.detach())
+            valid_block_counts.append(metrics.valid_blocks.detach())
+            loss_weights.append(metrics.loss_weight.detach())
             self.log("train/dflash_loss", metrics.loss, on_step=True, prog_bar=True)
             self.log("train/dflash_accuracy", metrics.accuracy, on_step=True)
             self.log("train/accept_len", metrics.accept_len, on_step=True)
@@ -624,9 +712,44 @@ def training_step(self, batch, batch_idx):
             for reason, count in skip_counts.items():
                 self.log(f"train/dflash_skip/{reason}", float(count), on_step=True)
             return torch.zeros((), device=self.device, requires_grad=True)
+        additive_metrics = {
+            "train/dflash_input_tokens": input_token_counts,
+            "train/dflash_valid_tokens": valid_token_counts,
+            "train/dflash_valid_blocks": valid_block_counts,
+            "train/dflash_loss_weight": loss_weights,
+        }
+        distributed_log_kwargs = {}
+        if self._draft_dp_size > 1:
+            distributed_log_kwargs = {
+                "sync_dist": True,
+                "sync_dist_group": self._draft_dp_group,
+                "reduce_fx": "sum",
+            }
+        for name, values in additive_metrics.items():
+            local_total = torch.stack([value.to(device=self.device, dtype=torch.float64) for value in values]).sum()
+            self.log(name, local_total, on_step=True, **distributed_log_kwargs)
         self.log("train/dflash_skipped_step", 0.0, on_step=True)
         return torch.stack(losses).mean()
 
+    def on_train_batch_end(self, outputs, batch, batch_idx) -> None:
+        """Persist post-backward peak HBM for production capacity checks."""
+        if not torch.cuda.is_available():
+            return
+        distributed_log_kwargs = {}
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            distributed_log_kwargs = {"sync_dist": True, "reduce_fx": "max"}
+        memory_metrics = {
+            "train/dflash_peak_memory_allocated_bytes": torch.cuda.max_memory_allocated(),
+            "train/dflash_peak_memory_reserved_bytes": torch.cuda.max_memory_reserved(),
+        }
+        for name, value in memory_metrics.items():
+            self.log(
+                name,
+                torch.tensor(value, device=self.device, dtype=torch.float64),
+                on_step=True,
+                **distributed_log_kwargs,
+            )
+
     def on_validation_epoch_start(self) -> None:
         self._partial_val_metrics.clear()
 
@@ -768,4 +891,7 @@ def on_train_end(self) -> None:
             return
         state_dict = _get_consolidated_model_state_dict(self.draft_model)
         if self.trainer.is_global_zero:
-            self.draft_model.save_pretrained(self.output_dir, state_dict=state_dict)
+            output_dir = Path(self.output_dir)
+            if not output_dir.is_absolute():
+                output_dir = Path(self.trainer.log_dir) / output_dir
+            self.draft_model.save_pretrained(output_dir, state_dict=state_dict)
diff --git a/nemo/collections/speechlm2/parts/packed_sequences.py b/nemo/collections/speechlm2/parts/packed_sequences.py
index 3df9b5ddad6b..7e211bd3fb8d 100644
--- a/nemo/collections/speechlm2/parts/packed_sequences.py
+++ b/nemo/collections/speechlm2/parts/packed_sequences.py
@@ -35,6 +35,231 @@
 from nemo.collections.speechlm2.parts.input_utils import _unpad_inputs
 
 
+def _validate_packed_dflash_inputs(packed: dict[str, Any]) -> None:
+    """Validate the token-aligned THD and DFlash document-boundary contract."""
+    required = {
+        "input_ids",
+        "input_embeddings",
+        "loss_mask",
+        "position_ids",
+        "seq_lens",
+        "doc_remaining",
+        "cu_seqlens",
+        "max_seqlen",
+        "qkv_format",
+    }
+    missing = sorted(required.difference(packed))
+    if missing:
+        raise ValueError(f"Packed DFlash inputs are missing required fields: {', '.join(missing)}")
+    if packed["qkv_format"] != "thd":
+        raise ValueError(f"Packed DFlash requires qkv_format='thd', got {packed['qkv_format']!r}")
+
+    input_ids = packed["input_ids"]
+    input_embeddings = packed["input_embeddings"]
+    loss_mask = packed["loss_mask"]
+    position_ids = packed["position_ids"]
+    seq_lens = packed["seq_lens"]
+    doc_remaining = packed["doc_remaining"]
+    cu_seqlens = packed["cu_seqlens"]
+    max_seqlen = packed["max_seqlen"]
+
+    if input_ids.ndim != 2 or input_ids.shape[0] != 1:
+        raise ValueError(f"Packed DFlash input_ids must have shape [1, T], got {tuple(input_ids.shape)}")
+    token_count = input_ids.shape[1]
+    if input_embeddings.ndim != 2 or input_embeddings.shape[0] != token_count:
+        raise ValueError(
+            "Packed DFlash input_embeddings must have shape [T, H] aligned with input_ids; "
+            f"got {tuple(input_embeddings.shape)} for T={token_count}"
+        )
+    for name, tensor in {
+        "loss_mask": loss_mask,
+        "position_ids": position_ids,
+        "doc_remaining": doc_remaining,
+    }.items():
+        if tensor.shape != input_ids.shape:
+            raise ValueError(
+                f"Packed DFlash {name} must have shape {tuple(input_ids.shape)}, got {tuple(tensor.shape)}"
+            )
+    if seq_lens.ndim != 2 or seq_lens.shape[0] != 1 or seq_lens.shape[1] == 0:
+        raise ValueError(f"Packed DFlash seq_lens must have shape [1, documents], got {tuple(seq_lens.shape)}")
+    if bool((seq_lens <= 0).any().item()) or int(seq_lens.sum().item()) != token_count:
+        raise ValueError(
+            "Packed DFlash document lengths must be positive and sum to the token count; "
+            f"got lengths={seq_lens.tolist()}, T={token_count}"
+        )
+
+    lengths = seq_lens.flatten().to(device=input_ids.device, dtype=torch.long)
+    expected_cu = torch.cat(
+        [
+            torch.zeros(1, dtype=torch.int32, device=input_ids.device),
+            lengths.cumsum(0, dtype=torch.int32),
+        ]
+    )
+    if not isinstance(cu_seqlens, Tensor) or cu_seqlens.dtype != torch.int32:
+        raise ValueError(
+            "Packed DFlash cu_seqlens must have dtype torch.int32, "
+            f"got {getattr(cu_seqlens, 'dtype', type(cu_seqlens).__name__)}"
+        )
+    if cu_seqlens.ndim != 1 or not torch.equal(cu_seqlens, expected_cu):
+        raise ValueError(
+            f"Packed DFlash cu_seqlens must equal cumulative seq_lens {expected_cu.tolist()}, "
+            f"got {cu_seqlens.tolist()}"
+        )
+    expected_max = int(lengths.max().item())
+    if not isinstance(max_seqlen, Tensor) or max_seqlen.numel() != 1 or int(max_seqlen.item()) != expected_max:
+        raise ValueError(f"Packed DFlash max_seqlen must be {expected_max}, got {max_seqlen}")
+
+    expected_positions = torch.cat(
+        [torch.arange(int(length), device=input_ids.device, dtype=torch.long) for length in lengths.tolist()]
+    ).unsqueeze(0)
+    expected_remaining = torch.cat(
+        [torch.arange(int(length) - 1, -1, -1, device=input_ids.device) for length in lengths.tolist()]
+    ).unsqueeze(0)
+    if not torch.equal(position_ids.to(expected_positions), expected_positions):
+        raise ValueError("Packed DFlash position_ids must reset to zero and increase within each document")
+    if not torch.equal(doc_remaining.to(expected_remaining), expected_remaining):
+        raise ValueError("Packed DFlash doc_remaining must count real tokens left within each document")
+
+
+def pack_audio_for_dflash(
+    input_ids: Tensor,
+    embeds: Tensor,
+    loss_mask: Tensor,
+    replacements: list[Tensor],
+    padding_id: int,
+    placeholder_id: int,
+    mask_token_id: int,
+) -> dict[str, Any]:
+    """Splice audio frames and concatenate utterances for packed SALM DFlash.
+
+    Unlike ordinary causal-LM packing, DFlash consumes the full unshifted token
+    stream. Each audio placeholder becomes one ``mask_token_id`` per encoder
+    frame and every inserted frame is excluded from supervision. The returned
+    target embeddings are flat THD ``[T, H]`` while all draft tensors retain a
+    synthetic packed batch row ``[1, T]``.
+    """
+    if input_ids.ndim != 2:
+        raise ValueError(f"Packed DFlash input_ids must have shape [B, S], got {tuple(input_ids.shape)}")
+    if embeds.ndim != 3 or embeds.shape[:2] != input_ids.shape:
+        raise ValueError(
+            "Packed DFlash embeds must have shape [B, S, H] matching input_ids; "
+            f"got {tuple(embeds.shape)} and {tuple(input_ids.shape)}"
+        )
+    if loss_mask.shape != input_ids.shape:
+        raise ValueError(
+            "Packed DFlash input_ids and loss_mask must have the same [B, S] shape; "
+            f"got {tuple(input_ids.shape)} and {tuple(loss_mask.shape)}"
+        )
+    if input_ids.shape[0] == 0:
+        raise ValueError("Packed DFlash requires at least one document")
+
+    hidden_size = embeds.shape[-1]
+    replacement_idx = 0
+    document_embeddings: list[Tensor] = []
+    document_ids: list[Tensor] = []
+    document_loss_masks: list[Tensor] = []
+    document_lengths: list[int] = []
+
+    for row_idx in range(input_ids.shape[0]):
+        row_ids = input_ids[row_idx]
+        non_padding = (row_ids != padding_id).nonzero(as_tuple=False)
+        if non_padding.numel() == 0:
+            raise ValueError(f"Packed DFlash document {row_idx} contains only padding")
+        start = int(non_padding[0].item())
+        row_ids = row_ids[start:]
+        row_embeddings = embeds[row_idx, start:]
+        row_loss_mask = loss_mask[row_idx, start:].bool()
+
+        embedding_segments: list[Tensor] = []
+        id_segments: list[Tensor] = []
+        loss_segments: list[Tensor] = []
+        previous = 0
+        for placeholder_pos in (row_ids == placeholder_id).nonzero(as_tuple=True)[0].tolist():
+            if placeholder_pos > previous:
+                ids_segment = row_ids[previous:placeholder_pos]
+                embedding_segments.append(row_embeddings[previous:placeholder_pos])
+                id_segments.append(ids_segment)
+                loss_segments.append(row_loss_mask[previous:placeholder_pos] & ids_segment.ne(padding_id))
+            if replacement_idx >= len(replacements):
+                raise ValueError("Packed DFlash has more audio placeholders than replacement tensors")
+            replacement = replacements[replacement_idx]
+            replacement_idx += 1
+            if replacement.ndim != 2 or replacement.shape[1] != hidden_size:
+                raise ValueError(
+                    "Packed DFlash audio replacements must have shape [frames, H] matching text embeddings; "
+                    f"got {tuple(replacement.shape)} for H={hidden_size}"
+                )
+            if replacement.shape[0] == 0:
+                raise ValueError("Packed DFlash audio replacements must contain at least one frame")
+            embedding_segments.append(replacement)
+            id_segments.append(
+                torch.full(
+                    (replacement.shape[0],),
+                    mask_token_id,
+                    dtype=input_ids.dtype,
+                    device=input_ids.device,
+                )
+            )
+            loss_segments.append(torch.zeros(replacement.shape[0], dtype=torch.bool, device=input_ids.device))
+            previous = placeholder_pos + 1
+
+        if previous < row_ids.numel():
+            ids_segment = row_ids[previous:]
+            embedding_segments.append(row_embeddings[previous:])
+            id_segments.append(ids_segment)
+            loss_segments.append(row_loss_mask[previous:] & ids_segment.ne(padding_id))
+        if not embedding_segments:
+            raise ValueError(f"Packed DFlash document {row_idx} has no real tokens after unpadding")
+
+        document_embedding = torch.cat(embedding_segments, dim=0)
+        document_id = torch.cat(id_segments, dim=0)
+        document_loss_mask = torch.cat(loss_segments, dim=0)
+        length = document_id.numel()
+        if document_embedding.shape[0] != length or document_loss_mask.numel() != length:
+            raise ValueError(f"Packed DFlash document {row_idx} is not token aligned after audio expansion")
+        document_embeddings.append(document_embedding)
+        document_ids.append(document_id)
+        document_loss_masks.append(document_loss_mask)
+        document_lengths.append(length)
+
+    if replacement_idx != len(replacements):
+        raise ValueError(
+            f"Packed DFlash used {replacement_idx} of {len(replacements)} audio replacements; "
+            "placeholder occurrences and replacements must match"
+        )
+
+    device = input_ids.device
+    flat_ids = torch.cat(document_ids).unsqueeze(0)
+    flat_loss_mask = torch.cat(document_loss_masks).unsqueeze(0)
+    seq_lens = torch.tensor(document_lengths, dtype=torch.long, device=device).unsqueeze(0)
+    position_ids = torch.cat(
+        [torch.arange(length, dtype=torch.long, device=device) for length in document_lengths]
+    ).unsqueeze(0)
+    doc_remaining = torch.cat(
+        [torch.arange(length - 1, -1, -1, dtype=torch.long, device=device) for length in document_lengths]
+    ).unsqueeze(0)
+    cu_seqlens = torch.cat(
+        [
+            torch.zeros(1, dtype=torch.int32, device=device),
+            seq_lens.flatten().cumsum(0, dtype=torch.int32),
+        ]
+    )
+    packed = {
+        "input_ids": flat_ids,
+        "input_embeddings": torch.cat(document_embeddings, dim=0),
+        "attention_mask": None,
+        "loss_mask": flat_loss_mask,
+        "position_ids": position_ids,
+        "seq_lens": seq_lens,
+        "doc_remaining": doc_remaining,
+        "cu_seqlens": cu_seqlens,
+        "max_seqlen": torch.tensor(max(document_lengths), dtype=torch.int32, device=device),
+        "qkv_format": "thd",
+    }
+    _validate_packed_dflash_inputs(packed)
+    return packed
+
+
 def pack_audio_into_text_embeds(
     input_ids: Tensor,
     embeds: Tensor,
diff --git a/tests/collections/speechlm2/test_salm_dflash.py b/tests/collections/speechlm2/test_salm_dflash.py
index e2ed7bcfbd37..cf27553458e9 100644
--- a/tests/collections/speechlm2/test_salm_dflash.py
+++ b/tests/collections/speechlm2/test_salm_dflash.py
@@ -26,12 +26,15 @@
 pytestmark = pytest.mark.unit
 
 from nemo_automodel.components.loss.dllm_loss import DFlashDecayLoss  # noqa: E402
-from nemo_automodel.components.speculative.dflash.draft_qwen3 import Qwen3DFlashDraftModel  # noqa: E402
+from nemo_automodel.components.speculative.dflash.draft_qwen3 import (
+    Qwen3DFlashDraftModel,
+)  # noqa: E402
 from nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2 import (  # noqa: E402
     Qwen3DFlash2DraftModel,
 )
 
 from nemo.collections.speechlm2.parts import dflash as salm_dflash  # noqa: E402
+from nemo.collections.speechlm2.parts import packed_sequences  # noqa: E402
 
 
 REPO_ROOT = Path(__file__).parents[3]
@@ -103,10 +106,13 @@ def test_synchronize_ep_group_is_noop_without_distributed_ep(monkeypatch):
     assert calls == []
 
 
-def test_validate_dflash_parallelism_rejects_tensor_parallelism():
-    mesh_context = SimpleNamespace(tp_size=2, pp_size=1, cp_size=1)
+@pytest.mark.parametrize("axis", ["tp_size", "cp_size"])
+def test_validate_dflash_parallelism_rejects_sequence_sharding(axis):
+    sizes = {"tp_size": 1, "pp_size": 1, "cp_size": 1}
+    sizes[axis] = 2
+    mesh_context = SimpleNamespace(**sizes)
 
-    with pytest.raises(NotImplementedError, match="tp_size=2"):
+    with pytest.raises(NotImplementedError, match=rf"{axis}=2"):
         salm_dflash._validate_dflash_parallelism(mesh_context)
 
 
@@ -242,6 +248,133 @@ def test_prepare_batch_keeps_full_unshifted_ids_and_token_aligned_loss_mask(
     assert module.trainer_module.kwargs["hidden_states"] is captured_hidden
 
 
+def test_pack_audio_for_dflash_builds_unshifted_boundary_metadata():
+    input_ids = torch.tensor([[0, 10, 99, 20, 21], [30, 31, 99, 0, 32]])
+    embeds = input_ids.to(torch.float32).unsqueeze(-1).expand(-1, -1, 2).clone()
+    loss_mask = torch.tensor(
+        [[False, False, False, True, True], [False, True, False, True, True]],
+        dtype=torch.bool,
+    )
+    replacements = [
+        torch.tensor([[100.0, 101.0], [102.0, 103.0]]),
+        torch.tensor([[200.0, 201.0]]),
+    ]
+
+    packed = packed_sequences.pack_audio_for_dflash(
+        input_ids=input_ids,
+        embeds=embeds,
+        loss_mask=loss_mask,
+        replacements=replacements,
+        padding_id=0,
+        placeholder_id=99,
+        mask_token_id=990,
+    )
+
+    assert packed["input_ids"].tolist() == [[10, 990, 990, 20, 21, 30, 31, 990, 0, 32]]
+    assert packed["loss_mask"].tolist() == [[False, False, False, True, True, False, True, False, False, True]]
+    assert packed["position_ids"].tolist() == [[0, 1, 2, 3, 4, 0, 1, 2, 3, 4]]
+    assert packed["seq_lens"].tolist() == [[5, 5]]
+    assert packed["doc_remaining"].tolist() == [[4, 3, 2, 1, 0, 4, 3, 2, 1, 0]]
+    assert packed["cu_seqlens"].dtype == torch.int32
+    assert packed["cu_seqlens"].tolist() == [0, 5, 10]
+    assert packed["max_seqlen"].item() == 5
+    assert packed["qkv_format"] == "thd"
+    assert packed["input_embeddings"].shape == (10, 2)
+    assert packed["input_embeddings"][1:3].tolist() == replacements[0].tolist()
+    assert packed["input_embeddings"][7:8].tolist() == replacements[1].tolist()
+
+
+def test_validate_packed_dflash_rejects_non_int32_cu_seqlens():
+    packed = {
+        "input_ids": torch.tensor([[10, 11, 20, 21]]),
+        "input_embeddings": torch.randn(4, 2),
+        "loss_mask": torch.ones(1, 4, dtype=torch.bool),
+        "position_ids": torch.tensor([[0, 1, 0, 1]]),
+        "seq_lens": torch.tensor([[2, 2]]),
+        "doc_remaining": torch.tensor([[1, 0, 1, 0]]),
+        "cu_seqlens": torch.tensor([0, 2, 4], dtype=torch.int64),
+        "max_seqlen": torch.tensor(2, dtype=torch.int32),
+        "qkv_format": "thd",
+    }
+
+    with pytest.raises(ValueError, match="cu_seqlens must have dtype torch.int32"):
+        packed_sequences._validate_packed_dflash_inputs(packed)
+
+
+def test_pack_audio_for_dflash_one_document_matches_unpacked_preparation(monkeypatch):
+    target = _BatchTarget()
+    module = salm_dflash.SALMDFlashModule(target, {"dflash": {"mask_token_id": 990, "block_size": 2}})
+    audio_embeddings = [torch.tensor([[100.0] * 4, [101.0] * 4])]
+    monkeypatch.setattr(module, "_audio_embeddings", Mock(return_value=audio_embeddings))
+    batch = {
+        "input_ids": torch.tensor([[0, 10, 99, 20, 21, 22]]),
+        "loss_mask": torch.tensor([[False, False, False, False, True, True]]),
+    }
+    unpacked = module._prepare_batch(batch)
+
+    target.cfg["packed_sequences"] = True
+    packed = module._prepare_batch(batch)
+
+    assert packed["input_ids"].tolist() == unpacked["input_ids"].tolist()
+    assert packed["loss_mask"].tolist() == unpacked["loss_mask"].tolist()
+    torch.testing.assert_close(packed["input_embeddings"].unsqueeze(0), unpacked["input_embeddings"])
+    assert packed["seq_lens"].tolist() == [[6]]
+    assert packed["doc_remaining"].tolist() == [[5, 4, 3, 2, 1, 0]]
+
+
+def test_pack_audio_for_dflash_rejects_malformed_inputs():
+    with pytest.raises(ValueError, match=r"same \[B, S\] shape"):
+        packed_sequences.pack_audio_for_dflash(
+            input_ids=torch.ones(1, 3, dtype=torch.long),
+            embeds=torch.ones(1, 3, 2),
+            loss_mask=torch.ones(1, 2, dtype=torch.bool),
+            replacements=[],
+            padding_id=0,
+            placeholder_id=99,
+            mask_token_id=990,
+        )
+
+
+@pytest.mark.parametrize("variant", ["dflash", "dflash2"])
+def test_run_batch_forwards_all_packing_metadata(monkeypatch, variant):
+    module = salm_dflash.SALMDFlashModule(
+        _BatchTarget(),
+        {"dflash": {"variant": variant, "mask_token_id": 990, "block_size": 2}},
+    )
+    packed = {
+        "input_ids": torch.tensor([[10, 11, 20, 21]]),
+        "input_embeddings": torch.randn(4, 4),
+        "attention_mask": None,
+        "loss_mask": torch.tensor([[True, True, True, True]]),
+        "position_ids": torch.tensor([[0, 1, 0, 1]]),
+        "seq_lens": torch.tensor([[2, 2]]),
+        "doc_remaining": torch.tensor([[1, 0, 1, 0]]),
+        "cu_seqlens": torch.tensor([0, 2, 4], dtype=torch.int32),
+        "max_seqlen": torch.tensor(2, dtype=torch.int32),
+        "qkv_format": "thd",
+    }
+    monkeypatch.setattr(module, "_prepare_batch", Mock(return_value=packed))
+    monkeypatch.setattr(module, "_target_hidden_states", Mock(return_value=torch.randn(1, 4, 8)))
+    monkeypatch.setattr(salm_dflash, "_has_valid_dflash_anchors", lambda *args, **kwargs: True)
+    module.trainer_module = _CaptureDFlashTrainer()
+    module._trainer = SimpleNamespace(strategy=SimpleNamespace(moe_mesh=None))
+
+    module._run_batch({})
+
+    assert module.trainer_module.kwargs["position_ids"] is packed["position_ids"]
+    assert module.trainer_module.kwargs["seq_lens"] is packed["seq_lens"]
+    assert module.trainer_module.kwargs["doc_remaining"] is packed["doc_remaining"]
+    assert module._last_input_token_count.item() == 4
+
+
+def test_packed_anchor_precheck_requires_complete_block_in_document():
+    loss_mask = torch.tensor([[False, True, True, True, True, True]])
+    doc_remaining = torch.tensor([[2, 1, 0, 2, 1, 0]])
+
+    assert not salm_dflash._has_valid_dflash_anchors(loss_mask, block_size=4, doc_remaining=doc_remaining)
+    assert salm_dflash._has_valid_dflash_anchors(loss_mask, block_size=3, doc_remaining=doc_remaining)
+
+
 def test_build_draft_config_applies_explicit_architecture_and_layer_taps():
     target_config = Qwen3Config(
         hidden_size=64,
@@ -325,7 +458,13 @@ def test_dflash2_rejects_fused_linear_ce():
     with pytest.raises(ValueError, match="use_fused_linear_ce"):
         salm_dflash.SALMDFlashModule(
             nn.Linear(1, 1),
-            {"dflash": {"variant": "dflash2", "mask_token_id": 18, "use_fused_linear_ce": True}},
+            {
+                "dflash": {
+                    "variant": "dflash2",
+                    "mask_token_id": 18,
+                    "use_fused_linear_ce": True,
+                }
+            },
         )
 
 
@@ -616,6 +755,149 @@ def test_target_hidden_states_filters_unsupported_optional_forward_kwargs():
     assert target.llm.calls == [{"attention_mask": inputs["attention_mask"]}]
 
 
+class _PackedTargetLLM(nn.Module):
+    """Small target that mixes causally inside, but never across, THD documents."""
+
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(1))
+        self.layers = nn.ModuleList([nn.Identity()])
+        self.calls = []
+
+    def forward(
+        self,
+        *,
+        inputs_embeds,
+        attention_mask=None,
+        qkv_format=None,
+        cu_seqlens=None,
+        position_ids=None,
+        max_seqlen=None,
+        output_hidden_states=False,
+        use_cache=False,
+        return_dict=True,
+        compute_logits=False,
+    ):
+        self.calls.append(
+            {
+                "qkv_format": qkv_format,
+                "cu_seqlens": cu_seqlens,
+                "position_ids": position_ids,
+                "max_seqlen": max_seqlen,
+                "compute_logits": compute_logits,
+            }
+        )
+        if qkv_format == "thd":
+            boundaries = cu_seqlens.tolist()
+            mixed = torch.cat(
+                [inputs_embeds[start:end].cumsum(dim=0) for start, end in zip(boundaries, boundaries[1:])],
+                dim=0,
+            )
+        else:
+            mixed = inputs_embeds.cumsum(dim=1)
+        hidden = self.layers[0](mixed)
+        return SimpleNamespace(hidden_states=(hidden,))
+
+
+class _PackedTargetModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.llm = _PackedTargetLLM()
+
+
+def _packed_target_inputs(embeddings, seq_lens):
+    lengths = torch.tensor([seq_lens], dtype=torch.long)
+    positions = torch.cat([torch.arange(length) for length in seq_lens]).unsqueeze(0)
+    remaining = torch.cat([torch.arange(length - 1, -1, -1) for length in seq_lens]).unsqueeze(0)
+    cu_seqlens = torch.tensor([0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32)
+    token_count = embeddings.shape[0]
+    return {
+        "input_ids": torch.arange(token_count).unsqueeze(0),
+        "input_embeddings": embeddings,
+        "attention_mask": None,
+        "loss_mask": torch.ones(1, token_count, dtype=torch.bool),
+        "position_ids": positions,
+        "seq_lens": lengths,
+        "doc_remaining": remaining,
+        "cu_seqlens": cu_seqlens,
+        "max_seqlen": torch.tensor(max(seq_lens), dtype=torch.int32),
+        "qkv_format": "thd",
+    }
+
+
+def test_target_hidden_states_one_document_packed_matches_unpacked():
+    target = _PackedTargetModel()
+    module = salm_dflash.SALMDFlashModule(target, {"dflash": {"mask_token_id": 18}})
+    module.target_layer_ids = [0]
+    embeddings = torch.randn(5, 4)
+
+    unpacked = module._target_hidden_states(
+        {
+            "input_embeddings": embeddings.unsqueeze(0),
+            "attention_mask": torch.ones(1, 5, dtype=torch.bool),
+        }
+    )
+    packed = module._target_hidden_states(_packed_target_inputs(embeddings, [5]))
+
+    torch.testing.assert_close(packed, unpacked)
+    assert packed.shape == (1, 5, 4)
+
+
+def test_target_hidden_states_packed_isolates_documents_and_uses_thd_metadata():
+    target = _PackedTargetModel()
+    module = salm_dflash.SALMDFlashModule(target, {"dflash": {"mask_token_id": 18}})
+    module.target_layer_ids = [0]
+    embeddings = torch.randn(6, 4)
+    inputs = _packed_target_inputs(embeddings, [3, 3])
+
+    reference = module._target_hidden_states(inputs)
+    independent = torch.cat(
+        [
+            module._target_hidden_states(
+                {
+                    "input_embeddings": embeddings[start:end].unsqueeze(0),
+                    "attention_mask": torch.ones(1, end - start, dtype=torch.bool),
+                }
+            )
+            for start, end in ((0, 3), (3, 6))
+        ],
+        dim=1,
+    )
+    torch.testing.assert_close(reference, independent)
+
+    perturbed_inputs = _packed_target_inputs(embeddings.clone(), [3, 3])
+    perturbed_inputs["input_embeddings"][3:] += 100
+    perturbed = module._target_hidden_states(perturbed_inputs)
+    torch.testing.assert_close(reference[:, :3], perturbed[:, :3])
+    assert not torch.allclose(reference[:, 3:], perturbed[:, 3:])
+
+    reverse_perturbed_inputs = _packed_target_inputs(embeddings.clone(), [3, 3])
+    reverse_perturbed_inputs["input_embeddings"][:3] += 100
+    reverse_perturbed = module._target_hidden_states(reverse_perturbed_inputs)
+    assert not torch.allclose(reference[:, :3], reverse_perturbed[:, :3])
+    torch.testing.assert_close(reference[:, 3:], reverse_perturbed[:, 3:])
+
+    assert target.llm.calls[-1]["qkv_format"] == "thd"
+    assert target.llm.calls[-1]["cu_seqlens"].tolist() == [0, 3, 6]
+    assert target.llm.calls[-1]["position_ids"].tolist() == [[0, 1, 2, 0, 1, 2]]
+    assert target.llm.calls[-1]["compute_logits"] is False
+
+
+def test_target_hidden_states_rejects_partial_packing_metadata():
+    module = salm_dflash.SALMDFlashModule(_PackedTargetModel(), {"dflash": {"mask_token_id": 18}})
+    module.target_layer_ids = [0]
+
+    with pytest.raises(ValueError, match="missing required fields"):
+        module._target_hidden_states(
+            {
+                "input_embeddings": torch.randn(4, 3),
+                "attention_mask": None,
+                "qkv_format": "thd",
+                "cu_seqlens": torch.tensor([0, 4], dtype=torch.int32),
+            }
+        )
+
+
 def test_get_consolidated_state_dict_uses_plain_state_dict_without_distributed(
     monkeypatch,
 ):
@@ -669,7 +951,13 @@ def fake_all_reduce(value, *, op, group):
 def test_dflash2_globally_normalizes_base_and_selector_terms_separately(monkeypatch):
     module = salm_dflash.SALMDFlashModule(
         nn.Linear(1, 1),
-        {"dflash": {"variant": "dflash2", "mask_token_id": 18, "selector_loss_weight": 0.5}},
+        {
+            "dflash": {
+                "variant": "dflash2",
+                "mask_token_id": 18,
+                "selector_loss_weight": 0.5,
+            }
+        },
     )
     module._draft_dp_size = 2
     module._draft_dp_group = object()
@@ -723,7 +1011,8 @@ def test_training_step_synchronizes_multi_dataset_skips(monkeypatch):
     module.draft_model = nn.Linear(1, 1)
     module._draft_dp_size = 1
     module._draft_dp_group = None
-    monkeypatch.setattr(module, "log", Mock())
+    log = Mock()
+    monkeypatch.setattr(module, "log", log)
     monkeypatch.setattr(salm_dflash, "_max_rank_value", lambda _value, _device: 3)
     availability = []
 
@@ -738,7 +1027,10 @@ def agree(local_condition, _device):
         loss_weight=torch.tensor(3.0),
         accuracy=torch.tensor(0.5),
         accept_len=torch.tensor(1.5),
+        valid_tokens=torch.tensor(12),
+        valid_blocks=torch.tensor(4),
     )
+    module._last_input_token_count = torch.tensor(21)
     run_batch = Mock(side_effect=[salm_dflash.NoValidAnchorsError("skip"), metrics])
     monkeypatch.setattr(module, "_run_batch", run_batch)
     batch = {
@@ -751,9 +1043,94 @@ def agree(local_condition, _device):
     torch.testing.assert_close(loss, metrics.loss)
     assert availability == [True, True, False]
     assert run_batch.call_count == 2
+    log.assert_any_call("train/dflash_input_tokens", torch.tensor(21.0, dtype=torch.float64), on_step=True)
+    log.assert_any_call("train/dflash_valid_tokens", torch.tensor(12.0, dtype=torch.float64), on_step=True)
+    log.assert_any_call("train/dflash_valid_blocks", torch.tensor(4.0, dtype=torch.float64), on_step=True)
+    log.assert_any_call("train/dflash_loss_weight", torch.tensor(3.0, dtype=torch.float64), on_step=True)
+
+
+def test_training_telemetry_sums_over_draft_dp_group(monkeypatch):
+    module = salm_dflash.SALMDFlashModule(nn.Linear(1, 1), {"dflash": {"mask_token_id": 18}})
+    module.draft_model = nn.Linear(1, 1)
+    module._draft_dp_size = 2
+    module._draft_dp_group = object()
+    module._last_input_token_count = torch.tensor(21)
+    log = Mock()
+    monkeypatch.setattr(module, "log", log)
+    monkeypatch.setattr(salm_dflash, "_max_rank_value", lambda value, _device: value)
+    monkeypatch.setattr(salm_dflash, "_all_ranks_agree", lambda condition, _device: condition)
+    monkeypatch.setattr(salm_dflash, "_all_ranks_report_same_value", lambda _value, _device: True)
+    metrics = SimpleNamespace(
+        loss=torch.tensor(2.0, requires_grad=True),
+        loss_weight=torch.tensor(3.0),
+        accuracy=torch.tensor(0.5),
+        accept_len=torch.tensor(1.5),
+        valid_tokens=torch.tensor(12),
+        valid_blocks=torch.tensor(4),
+    )
+    monkeypatch.setattr(module, "_run_batch", Mock(return_value=metrics))
+    monkeypatch.setattr(module, "_globally_normalized_loss", lambda result: result.loss)
+
+    module.training_step({"input_ids": torch.ones(1, 2, dtype=torch.long)}, batch_idx=0)
+
+    log.assert_any_call(
+        "train/dflash_input_tokens",
+        torch.tensor(21.0, dtype=torch.float64),
+        on_step=True,
+        sync_dist=True,
+        sync_dist_group=module._draft_dp_group,
+        reduce_fx="sum",
+    )
+
+
+def test_training_peak_memory_telemetry_uses_max_rank_value(monkeypatch):
+    module = salm_dflash.SALMDFlashModule(nn.Linear(1, 1), {"dflash": {"mask_token_id": 18}})
+    log = Mock()
+    monkeypatch.setattr(module, "log", log)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: 123)
+    monkeypatch.setattr(torch.cuda, "max_memory_reserved", lambda: 456)
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+
+    module.on_train_batch_end(outputs=None, batch=None, batch_idx=0)
+
+    log.assert_any_call(
+        "train/dflash_peak_memory_allocated_bytes",
+        torch.tensor(123.0, dtype=torch.float64),
+        on_step=True,
+        sync_dist=True,
+        reduce_fx="max",
+    )
+    log.assert_any_call(
+        "train/dflash_peak_memory_reserved_bytes",
+        torch.tensor(456.0, dtype=torch.float64),
+        on_step=True,
+        sync_dist=True,
+        reduce_fx="max",
+    )
+
 
+@pytest.mark.parametrize("configured", [Path("outputs/draft"), Path("/durable/draft")])
+def test_train_end_resolves_relative_export_under_log_dir(monkeypatch, tmp_path, configured):
+    module = salm_dflash.SALMDFlashModule(
+        nn.Linear(1, 1),
+        {"dflash": {"mask_token_id": 18, "output_dir": str(configured)}},
+    )
+    module.draft_model = Mock()
+    module._trainer = SimpleNamespace(is_global_zero=True, log_dir=str(tmp_path / "experiment"))
+    state_dict = {"weight": torch.tensor([1.0])}
+    monkeypatch.setattr(salm_dflash, "_get_consolidated_model_state_dict", lambda _model: state_dict)
 
-def test_training_step_returns_differentiable_zero_when_every_dataset_is_skipped(monkeypatch):
+    module.on_train_end()
+
+    expected = configured if configured.is_absolute() else tmp_path / "experiment" / configured
+    module.draft_model.save_pretrained.assert_called_once_with(expected, state_dict=state_dict)
+
+
+def test_training_step_returns_differentiable_zero_when_every_dataset_is_skipped(
+    monkeypatch,
+):
     module = salm_dflash.SALMDFlashModule(nn.Linear(1, 1), {"dflash": {"mask_token_id": 18}})
     module.draft_model = nn.Linear(1, 1)
     log = Mock()
@@ -814,10 +1191,18 @@ def test_aggregate_validation_accuracy_preserves_default_checkpoint_monitor(
     log.assert_any_call("val_acc", torch.tensor(0.5), on_epoch=True)
 
 
-def test_dflash2_validation_uses_separate_loss_denominators_and_selector_metrics(monkeypatch):
+def test_dflash2_validation_uses_separate_loss_denominators_and_selector_metrics(
+    monkeypatch,
+):
     module = salm_dflash.SALMDFlashModule(
         nn.Linear(1, 1),
-        {"dflash": {"variant": "dflash2", "mask_token_id": 18, "selector_loss_weight": 0.5}},
+        {
+            "dflash": {
+                "variant": "dflash2",
+                "mask_token_id": 18,
+                "selector_loss_weight": 0.5,
+            }
+        },
     )
     log = Mock()
     monkeypatch.setattr(module, "log", log)