feat(speechlm2): add DFlash and DFlash2 training for SALM

docs/source/speechlm2/intro.rst

@@ -370,6 +370,22 @@ 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 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>`.
 
 Collection Structure

examples/speechlm2/conf/salm_automodel.yaml

@@ -290,3 +290,57 @@ 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.
+# 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 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]
+  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: 512
+  max_total_anchors: 512
+  loss_decay_gamma: 4.0
+  draft_sliding_window: null
+  attention_backend: flex_attention
+  activation_checkpointing: 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
+  output_dir: ./outputs/salm_dflash

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,25 @@ 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)
+        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
+
     model_cls = SALM
     if cfg.model.get("use_nemo_automodel", False):
         from nemo.collections.speechlm2 import SALMAutomodel