feat(dllm): add Transformer Engine context parallelism for DiffusionGemma

examples/dllm_sft/dflash_sft.yaml

@@ -49,7 +49,7 @@ model:
   _target_: transformers.AutoModel.from_pretrained
   pretrained_model_name_or_path: z-lab/Qwen3-4B-DFlash-b16
   trust_remote_code: true
-  dtype: bfloat16
+  dtype: float32
 
 checkpoint:
   enabled: true
@@ -65,7 +65,7 @@ distributed:
   sequence_parallel: false
   activation_checkpointing: false
   mp_policy:
-    param_dtype: bfloat16
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16

examples/dllm_sft/diffusion_gemma_lora.yaml

@@ -51,7 +51,7 @@ wandb:
 model:
   _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
   pretrained_model_name_or_path: google/diffusiongemma-26B-A4B-it
-  torch_dtype: float32        # fp32 master weights; compute is bf16 via mp_policy
+  torch_dtype: float32        # fp32 master weights; compute is bf16 via autocast
   canvas_length: 256
   self_conditioning: true
   freeze_router: true
@@ -100,8 +100,8 @@ distributed:
   ep_size: 8
   sequence_parallel: false
   activation_checkpointing: true
-  mp_policy:                  # mixed precision: fp32 master + bf16 compute
-    param_dtype: bfloat16
+  mp_policy:                  # mixed precision: fp32 master + bf16 autocast compute
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16

examples/dllm_sft/diffusion_gemma_sft.yaml

@@ -54,7 +54,7 @@ wandb:
 model:
   _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
   pretrained_model_name_or_path: google/diffusiongemma-26B-A4B-it
-  torch_dtype: float32        # fp32 master weights; compute is bf16 via mp_policy
+  torch_dtype: float32        # fp32 master weights; compute is bf16 via autocast
   canvas_length: 256
   self_conditioning: true
   freeze_router: true
@@ -86,8 +86,8 @@ distributed:
   ep_size: 8
   sequence_parallel: false
   activation_checkpointing: true
-  mp_policy:                  # mixed precision: fp32 master + bf16 compute
-    param_dtype: bfloat16
+  mp_policy:                  # mixed precision: fp32 master + bf16 autocast compute
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16

examples/dllm_sft/diffusion_gemma_te_cp_100k.yaml

@@ -0,0 +1,115 @@
+# Two-node (16 GPU) DiffusionGemma SFT proof at a genuine 100K-token context.
+#
+# Prepare four streamed FineWeb-Edu passkey samples:
+#   python examples/dllm_sft/prep_diffusion_gemma_long_context.py \
+#     --target-tokens 100000 --num-samples 4
+#
+# Launch on two 8-GPU nodes with torchrun. CP=16 is the entire multi-node world.
+
+recipe: DiffusionGemmaSFTRecipe
+seed: 42
+
+step_scheduler:
+  global_batch_size: 1
+  local_batch_size: 1
+  log_remote_every_steps: 1
+  ckpt_every_steps: 100000
+  val_every_steps: 100000
+  max_steps: 2
+  num_epochs: 1
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 30
+
+model:
+  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
+  pretrained_model_name_or_path: google/diffusiongemma-26B-A4B-it
+  torch_dtype: float32
+  canvas_length: 256
+  self_conditioning: true
+  freeze_router: true
+  backend:
+    _target_: nemo_automodel.components.models.common.BackendConfig
+    attn: te
+    linear: torch
+    rms_norm: torch_fp32
+    experts: torch_mm
+    dispatcher: torch
+    enable_hf_state_dict_adapter: true
+    enable_fsdp_optimizations: true
+
+distributed:
+  strategy: fsdp2
+  dp_size: none
+  tp_size: 1
+  cp_size: 16
+  pp_size: 1
+  ep_size: 1
+  sequence_parallel: false
+  activation_checkpointing: true
+  mp_policy:
+    param_dtype: float32
+    reduce_dtype: float32
+    output_dtype: float32
+  autocast_dtype: bfloat16
+  offload_policy: null
+  moe:
+    reshard_after_forward: false
+
+dllm:
+  mode: block_diffusion
+  block_size: 256
+  vocab_size: 262144
+  eps: 0.001
+  pad_block_size: 256
+  pad_seq_len_divisible: 256
+  encoder_loss_weight: 1.0
+
+loss_fn:
+  _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy
+
+optimizer:
+  _target_: torch.optim.AdamW
+  betas: [0.95, 0.99]
+  eps: 1.0e-8
+  lr: 1.5e-4
+  weight_decay: 1.0e-4
+
+clip_grad_norm:
+  max_norm: 1.0
+
+lr_scheduler:
+  lr_warmup_steps: 1
+  init_lr: 0.0
+  lr_decay_style: cosine
+  min_lr: 1.5e-5
+
+checkpoint:
+  enabled: false
+
+dataset:
+  _target_: nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset
+  path_or_dataset_id: diffusion_gemma_long_100k.jsonl
+  split: train
+  shuffle_seed: 42
+  seq_length: 100096
+  truncation: true
+  unshifted: true
+  mask_history: true
+  chat_template: examples/dllm_sft/diffusion_gemma_chat_template.jinja
+  tokenizer:
+    pretrained_model_name_or_path: google/diffusiongemma-26B-A4B-it
+
+dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  collate_fn: nemo_automodel.components.datasets.utils.default_collater
+  shuffle: true
+
+wandb:
+  enable: false
+
+ci:
+  recipe_owner: akoumpa
+  nodes: 2
+  time: "02:00:00"

examples/dllm_sft/llada2_lora.yaml

@@ -66,7 +66,7 @@ wandb:
 model:
   _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
   pretrained_model_name_or_path: inclusionAI/LLaDA2.1-mini
-  torch_dtype: bfloat16
+  torch_dtype: float32
   trust_remote_code: true
 
 checkpoint:
@@ -97,7 +97,7 @@ distributed:
   sequence_parallel: false
   activation_checkpointing: false
   mp_policy:
-    param_dtype: bfloat16
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16

examples/dllm_sft/llada2_sft.yaml

@@ -55,7 +55,7 @@ wandb:
 model:
   _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
   pretrained_model_name_or_path: inclusionAI/LLaDA2.1-mini
-  torch_dtype: bfloat16
+  torch_dtype: float32
   trust_remote_code: true
 
 checkpoint:
@@ -72,7 +72,7 @@ distributed:
   sequence_parallel: false
   activation_checkpointing: false
   mp_policy:
-    param_dtype: bfloat16
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16

examples/dllm_sft/llada_lora.yaml

@@ -78,7 +78,7 @@ distributed:
   sequence_parallel: false
   activation_checkpointing: false
   mp_policy:
-    param_dtype: bfloat16
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16

examples/dllm_sft/llada_scdd.yaml

@@ -76,7 +76,7 @@ wandb:
 model:
   _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
   pretrained_model_name_or_path: GSAI-ML/LLaDA-8B-Base
-  torch_dtype: float32     # fp32 master weights; compute stays bf16 via mp_policy
+  torch_dtype: float32     # fp32 master weights; compute stays bf16 via autocast
   trust_remote_code: true
 
 checkpoint:
@@ -94,7 +94,7 @@ distributed:
   sequence_parallel: false
   activation_checkpointing: true   # 8B full FT; trade compute for memory
   mp_policy:
-    param_dtype: bfloat16
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16

examples/dllm_sft/llada_sft.yaml

@@ -54,7 +54,7 @@ distributed:
   sequence_parallel: false
   activation_checkpointing: false
   mp_policy:
-    param_dtype: bfloat16
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16

examples/dllm_sft/nemotron_labs_diffusion_lora.yaml

@@ -82,7 +82,7 @@ distributed:
   sequence_parallel: false
   activation_checkpointing: true
   mp_policy:
-    param_dtype: bfloat16
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16

examples/dllm_sft/nemotron_labs_diffusion_sft.yaml

@@ -59,7 +59,7 @@ distributed:
   sequence_parallel: false
   activation_checkpointing: true
   mp_policy:
-    param_dtype: bfloat16
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16

examples/dllm_sft/nemotron_nano30b_dflash.yaml

@@ -56,7 +56,7 @@ model:
   _target_: transformers.AutoModel.from_pretrained
   pretrained_model_name_or_path: checkpoints/nemotron-nano-30b-dflash-b16/init
   trust_remote_code: true
-  dtype: bfloat16
+  dtype: float32
 
 checkpoint:
   enabled: true
@@ -72,7 +72,7 @@ distributed:
   sequence_parallel: false
   activation_checkpointing: false
   mp_policy:
-    param_dtype: bfloat16
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16

examples/dllm_sft/prep_diffusion_gemma_long_context.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Build a tiny, real-text 100K--256K DiffusionGemma SFT dataset.
+
+The source text is streamed from FineWeb-Edu, so only enough documents for the
+requested examples are downloaded. Each example is a single-turn passkey task:
+the long document is context, while the supervised response stays one short
+DiffusionGemma canvas.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+
+def _render_length(tokenizer, user: str, answer: str) -> int:
+    messages = [{"role": "user", "content": user}, {"role": "assistant", "content": answer}]
+    return len(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False))
+
+
+def main() -> None:
+    """Stream FineWeb-Edu and write exact-length passkey chat examples."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--output", default="diffusion_gemma_long_100k.jsonl")
+    parser.add_argument("--target-tokens", type=int, default=100_000, choices=(100_000, 131_072, 256_000))
+    parser.add_argument("--num-samples", type=int, default=4)
+    parser.add_argument("--model", default="google/diffusiongemma-26B-A4B-it")
+    args = parser.parse_args()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    documents = iter(load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT", split="train", streaming=True))
+
+    with open(args.output, "w", encoding="utf-8") as output:
+        for sample_idx in range(args.num_samples):
+            passkey = f"DG-LONG-{sample_idx:04d}-7391"
+            prefix = (
+                f"Remember this passkey: {passkey}. Read the document below. "
+                "At the end, return only the passkey.\n\nDOCUMENT START\n"
+            )
+            suffix = "\nDOCUMENT END\nReturn only the passkey."
+            answer = passkey
+
+            corpus_tokens: list[int] = []
+            # Collect with margin for the chat template and instructions.
+            while len(corpus_tokens) < args.target_tokens:
+                text = next(documents)["text"]
+                corpus_tokens.extend(tokenizer.encode(text + "\n\n", add_special_tokens=False))
+
+            lo, hi = 0, min(len(corpus_tokens), args.target_tokens)
+            best_user, best_length = "", 0
+            while lo <= hi:
+                mid = (lo + hi) // 2
+                corpus = tokenizer.decode(corpus_tokens[:mid], skip_special_tokens=True)
+                user = prefix + corpus + suffix
+                length = _render_length(tokenizer, user, answer)
+                if length <= args.target_tokens:
+                    best_user, best_length = user, length
+                    lo = mid + 1
+                else:
+                    hi = mid - 1
+
+            row = {
+                "messages": [
+                    {"role": "user", "content": best_user},
+                    {"role": "assistant", "content": answer},
+                ],
+                "source": "HuggingFaceFW/fineweb-edu/sample-10BT",
+                "rendered_tokens": best_length,
+            }
+            output.write(json.dumps(row, ensure_ascii=False) + "\n")
+            print(f"sample={sample_idx} rendered_tokens={best_length} passkey={passkey}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

examples/dllm_sft/qwen3_4b_dflash.yaml

@@ -50,7 +50,7 @@ model:
   _target_: transformers.AutoModel.from_pretrained
   pretrained_model_name_or_path: z-lab/Qwen3-4B-DFlash-b16
   trust_remote_code: true
-  dtype: bfloat16
+  dtype: float32
 
 checkpoint:
   enabled: true
@@ -66,7 +66,7 @@ distributed:
   sequence_parallel: false
   activation_checkpointing: false
   mp_policy:
-    param_dtype: bfloat16
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16

examples/dllm_sft/qwen3_8b_idlm.yaml

@@ -53,7 +53,7 @@ model:
   _target_: transformers.AutoModelForCausalLM.from_pretrained
   pretrained_model_name_or_path: Qwen/Qwen3-8B
   trust_remote_code: true
-  dtype: bfloat16
+  dtype: float32
   # I-DLM feeds a custom [x_t | x_0] block-diffusion mask; sdpa honours it.
   # Use flex_attention to avoid materialising the dense 2L x 2L mask at scale.
   attn_implementation: sdpa
@@ -74,7 +74,7 @@ distributed:
   sequence_parallel: false
   activation_checkpointing: true   # 8B full FT at seq 4096; trade compute for memory
   mp_policy:
-    param_dtype: bfloat16
+    param_dtype: float32
     reduce_dtype: float32
     output_dtype: float32
   autocast_dtype: bfloat16