feat(dllm): add AutoModel backend for BlockJustGRPO

examples/configs/grpo_sudoku6x6_block_just_grpo_fast_1n.yaml

@@ -0,0 +1,56 @@
+# BlockJustGRPO-Fast, 6x6 Sudoku, SGLang, SINGLE NODE (1x8 GPUs).
+#
+# Why this file exists: no 1-node Fast config ships with the fork.
+#   - `..._toy_8p8g_fast.yaml`  is Fast but inherits cluster 16x8 / max_new_tokens 3744
+#   - `..._toy_8p8g_1n.yaml`    is 1 node / 512 tokens but NOT Fast
+# They cannot simply be combined in one `defaults:` list, because the loader fully
+# resolves each parent before merging, so the Fast parent's num_nodes=16 wins. This
+# child therefore inherits the 1-node config and re-adds only the Fast knobs.
+#
+# Dataset: sudoku6x6 is generated locally by
+# nemo_rl/data/datasets/response_datasets/sudoku6x6_generator.py -- no HF download,
+# so the job needs no outbound network at runtime.
+#
+# THREE PLACES block_size MUST AGREE (a mismatch does not error; it silently
+# produces a wrong objective -- the fork records gen_kl 1.41 from exactly this):
+#   1. policy.hf_config_overrides.block_size .............. 16  (inherited)
+#   2. dllm_algorithm_config -> ..._bs16_entropy.yaml ..... block_size: 16, max_steps: 16
+#   3. policy.logprob_estimation.block_size ............... OMITTED, so it inherits (1)
+#   + policy.generation.max_new_tokens must be a multiple of 16 (512 is)
+#
+# Fast requires per-token entropy from the generation engine. Only the SGLang fork
+# TIP (restore/a652eb48-good-run @ c02140e) emits `output_token_entropy_val`; the
+# commit the fork's pyproject originally pinned (a652eb48) has no entropy code at
+# all, so BlockJustGRPO-Fast is impossible on it. If the channel is missing the run
+# fail-fasts at require_generation_entropy rather than silently sparsifying on zeros.
+defaults:
+  - "deepscaler_nemotron_labs_diffusion_3b_sglang_block_just_grpo_megatron_toy_8p8g_1n.yaml"
+  - "grpo_sudoku6x6_base.yaml"
+
+policy:
+  logprob_estimation:
+    # m = ceil(0.25 * 16) = 4 reveal levels / forward passes instead of 16.
+    fast_entropy_level_ratio: 0.25
+  generation:
+    max_new_tokens: 512   # multiple of block_size 16
+    sglang_cfg:
+      # In-repo relative path. The shipped configs point into /home/snorouzi/... ,
+      # which is unreadable to us; the identical file is tracked here.
+      dllm_algorithm_config: "tools/nemotron_diffusion/block_just_grpo_leftmost_dllm_bs16_entropy.yaml"
+
+cluster:
+  num_nodes: 1
+  gpus_per_node: 8
+
+grpo:
+  max_num_steps: 2
+  val_at_start: false
+  val_at_end: false
+
+# Compatibility probe, not an experiment: no checkpoints, no W&B.
+checkpointing:
+  enabled: false
+
+logger:
+  wandb_enabled: false
+  log_dir: "logs/bjgf-sudoku6x6-1n"

examples/configs/grpo_sudoku6x6_block_just_grpo_fast_1n_automodel.yaml

@@ -0,0 +1,89 @@
+# BlockJustGRPO-Fast on 6x6 Sudoku, 1 node x 8 GPUs, AUTOMODEL (DTensor) backend.
+#
+# Backend-swapped sibling of grpo_sudoku6x6_block_just_grpo_fast_1n.yaml, which runs
+# the same algorithm on Megatron and is the parity target (that config reached
+# reward 0.110 -> 0.307 over 30 steps at gen-KL ~1.8e-3).
+#
+# Everything above the policy worker is inherited unchanged: the GRPO loop, the
+# dataloader, the reasoning_gym environment, advantages, and the SGLang FastDiffuser
+# generation path. Only the trainer differs.
+defaults: "grpo_sudoku6x6_block_just_grpo_fast_1n.yaml"
+
+policy:
+  # --- the backend swap -----------------------------------------------------
+  worker_cls_fqn: "nemo_rl.models.policy.workers.block_just_grpo_dtensor_policy_worker.BlockJustGRPODTensorPolicyWorker"
+
+  megatron_cfg:
+    enabled: false
+
+  dtensor_cfg:
+    enabled: true
+    _v2: true                    # NOT the code default; selects the Automodel worker
+    tensor_parallel_size: 1
+    context_parallel_size: 1     # CP is unsupported on this path -- see the worker's guard
+    automodel_kwargs:
+      # NLD is loaded through its HF remote-code implementation. AutoModel still
+      # supplies FSDP2/optimizer/refit around that model.
+      force_hf: true
+    # The Megatron run had activation checkpointing ON. Fast runs 4 reveal-level
+    # forwards per step over a 2x-length [noisy|clean] input, so memory pressure is
+    # higher here than for an ordinary LLM. Turn on if it OOMs.
+    activation_checkpointing: false
+
+  # Automodel builds its optimizer from policy.optimizer; the Megatron recipes set
+  # this to null and configure megatron_cfg.optimizer instead. The two schemas are
+  # NOT 1:1, so this has to be restated rather than translated mechanically.
+  optimizer:
+    name: "torch.optim.AdamW"
+    kwargs:
+      lr: 3.0e-7
+      weight_decay: 0.01
+      betas: [0.9, 0.999]
+      eps: 1.0e-8
+      # Required for DTensor/FSDP parameters; these match the stock AutoModel
+      # GRPO recipe rather than inheriting Megatron's fused optimizer path.
+      foreach: false
+      fused: false
+
+  # The Megatron parent sets policy.scheduler=null because mcore owns scheduling.
+  # AutoModel owns it here. Match the Megatron recipe's 13-step warmup from
+  # 3e-8 to 3e-7, followed by a constant learning rate.
+  scheduler:
+    - name: "torch.optim.lr_scheduler.LinearLR"
+      kwargs:
+        start_factor: 0.1
+        end_factor: 1.0
+        total_iters: 13
+    - name: "torch.optim.lr_scheduler.ConstantLR"
+      kwargs:
+        factor: 1.0
+        total_iters: 10000000000
+    - milestones: [13]
+
+  # --- the model ------------------------------------------------------------
+  # Patched clone: weights symlinked, code fixed for transformers 5.3.0, and
+  # auto_map extended with AutoModelForCausalLM. See run/patch_checkpoint.py.
+  model_name: "/lustre/fsw/portfolios/coreai/users/zezhou/dev/diffusionllm/dllm-rl-integration/build/checkpoints/Nemotron-Labs-Diffusion-3B"
+
+  hf_config_overrides:
+    block_size: 16
+    # CRITICAL, and verified on hardware. MinistralDiffEncoderModel.__init__ picks
+    # its attention class from dlm_paradigm:
+    #     block_diff / sbd_block_diff  -> MinistralFlexAttention (asymmetric mask)
+    #     bidirectional / autoregressive -> plain causal attention
+    # The checkpoint ships 'autoregressive'. Without this override the run trains
+    # happily on ORDINARY CAUSAL ATTENTION and is silently wrong. The worker also
+    # asserts this at startup.
+    dlm_paradigm: "sbd_block_diff"
+
+  # Packing keeps documents apart via cu_seqlens because attention is causal;
+  # bidirectional attention inside a packed buffer leaks across boundaries.
+  sequence_packing:
+    enabled: false
+
+checkpointing:
+  enabled: false
+
+logger:
+  wandb_enabled: false
+  log_dir: "logs/bjgf-sudoku6x6-1n-automodel"

nemo_rl/distributed/model_utils.py

@@ -13,19 +13,16 @@
 # limitations under the License.
 
 from collections.abc import Sequence
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import torch
-from megatron.core.models.gpt import GPTModel
-from megatron.core.parallel_state import (
-    get_context_parallel_group,
-    get_context_parallel_world_size,
-    get_tensor_model_parallel_group,
-    get_tensor_model_parallel_rank,
-)
-from megatron.core.utils import deprecate_inference_params, get_pg_size
 from torch.distributed.tensor import DTensor, distribute_tensor
 
+if TYPE_CHECKING:
+    # megatron-core (optional "mcore" extra) is imported lazily inside the
+    # functions below so this module imports without mcore installed.
+    from megatron.core.models.gpt import GPTModel
+
 from nemo_rl.algorithms.logits_sampling_utils import (
     TrainingSamplingParams,
     apply_top_k_top_p,
@@ -1071,12 +1068,20 @@ def from_parallel_logits_to_same_position_logprobs(
         )
 
     batch_size, seq_len, _ = vocab_parallel_logits.shape
-    target_positions = target_positions.to(device=vocab_parallel_logits.device, dtype=torch.long)
-    target_tokens = target_tokens.to(device=vocab_parallel_logits.device, dtype=torch.long)
+    target_positions = target_positions.to(
+        device=vocab_parallel_logits.device, dtype=torch.long
+    )
+    target_tokens = target_tokens.to(
+        device=vocab_parallel_logits.device, dtype=torch.long
+    )
     if torch.any((target_positions < 0) | (target_positions >= seq_len)):
-        raise ValueError(f"target_positions must be in [0, {seq_len}), got {target_positions}")
+        raise ValueError(
+            f"target_positions must be in [0, {seq_len}), got {target_positions}"
+        )
 
-    gather_positions = (target_positions + int(position_shift)).clamp(min=0, max=seq_len - 1)
+    gather_positions = (target_positions + int(position_shift)).clamp(
+        min=0, max=seq_len - 1
+    )
     row_indices = torch.arange(batch_size, device=vocab_parallel_logits.device)
     selected_logits = vocab_parallel_logits[row_indices, gather_positions, :]
     if (
@@ -1313,6 +1318,11 @@ def gather_cp_sharded_logits(output_tensor: torch.Tensor) -> torch.Tensor:
     logits are reconstructed (seq_dim=1) first. No-op when cp_size <= 1. Shared
     by the JustGRPO and DiffuGRPO Megatron post-processors.
     """
+    from megatron.core.parallel_state import (
+        get_context_parallel_group,
+        get_context_parallel_world_size,
+    )
+
     cp_size = get_context_parallel_world_size()
     if cp_size <= 1:
         return output_tensor
@@ -2166,6 +2176,8 @@ def backward(
 
 
 def patch_gpt_model_forward_for_linear_ce_fusion(*, chunk_size: int) -> None:
+    from megatron.core.models.gpt import GPTModel
+
     if getattr(GPTModel, "_linear_ce_fusion_forward_patched", False):
         GPTModel._linear_ce_fusion_chunk_size = chunk_size
         return
@@ -2176,7 +2188,7 @@ def patch_gpt_model_forward_for_linear_ce_fusion(*, chunk_size: int) -> None:
 
 
 def _gpt_forward_with_linear_ce_fusion(
-    self: GPTModel,
+    self: "GPTModel",
     input_ids: torch.Tensor,
     position_ids: torch.Tensor,
     attention_mask: torch.Tensor,
@@ -2192,6 +2204,12 @@ def _gpt_forward_with_linear_ce_fusion(
     padding_mask: Optional[torch.Tensor] = None,
     return_logprobs_for_linear_ce_fusion: bool = False,
 ) -> torch.Tensor:
+    from megatron.core.parallel_state import (
+        get_tensor_model_parallel_group,
+        get_tensor_model_parallel_rank,
+    )
+    from megatron.core.utils import deprecate_inference_params, get_pg_size
+
     if not return_logprobs_for_linear_ce_fusion:
         return self._original_forward_for_linear_ce_fusion(
             input_ids=input_ids,

nemo_rl/distributed/ray_actor_environment_registry.py

@@ -31,13 +31,17 @@
 MCORE_EXECUTABLE = (
     PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.MCORE
 )
+AUTOMODEL_EXECUTABLE = os.environ.get("NRL_AUTOMODEL_PY_EXECUTABLE") or (
+    PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.AUTOMODEL
+)
 
 ACTOR_ENVIRONMENT_REGISTRY: dict[str, str] = {
     "nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker": VLLM_EXECUTABLE,
     "nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker": VLLM_EXECUTABLE,
     "nemo_rl.models.generation.sglang.sglang_worker.SGLangGenerationWorker": SGLANG_EXECUTABLE,
     "nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker": PY_EXECUTABLES.FSDP,
-    "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2": PY_EXECUTABLES.AUTOMODEL,
+    "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2": AUTOMODEL_EXECUTABLE,
+    "nemo_rl.models.policy.workers.block_just_grpo_dtensor_policy_worker.BlockJustGRPODTensorPolicyWorker": AUTOMODEL_EXECUTABLE,
     "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker": MCORE_EXECUTABLE,
     "nemo_rl.models.policy.workers.nemotron_diffusion_megatron_policy_worker.NemotronDiffusionMegatronPolicyWorker": MCORE_EXECUTABLE,
     "nemo_rl.models.policy.workers.just_grpo_megatron_policy_worker.JustGRPOMegatronPolicyWorker": MCORE_EXECUTABLE,

nemo_rl/models/automodel/data.py

@@ -51,6 +51,11 @@ class ProcessedInputs:
     # Multimodal (VLM) inputs
     vlm_kwargs: dict[str, Any] = field(default_factory=dict)
 
+    # Optional model-specific arguments. This is deliberately separate from
+    # ``vlm_kwargs``: diffusion language models, for example, need to pass the
+    # denoising block size without pretending the input is multimodal.
+    model_kwargs: dict[str, Any] = field(default_factory=dict)
+
     # Context parallel support (cp_size > 1)
     cp_buffers: list[torch.Tensor] = field(default_factory=list)
     seq_index: Optional[torch.Tensor] = None