diff --git a/examples/configs/grpo_sudoku6x6_block_just_grpo_fast_1n.yaml b/examples/configs/grpo_sudoku6x6_block_just_grpo_fast_1n.yaml new file mode 100644 index 00000000000..bfd5bf89f3c --- /dev/null +++ b/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" diff --git a/examples/configs/grpo_sudoku6x6_block_just_grpo_fast_1n_automodel.yaml b/examples/configs/grpo_sudoku6x6_block_just_grpo_fast_1n_automodel.yaml new file mode 100644 index 00000000000..9a9e4937e17 --- /dev/null +++ b/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" diff --git a/nemo_rl/distributed/model_utils.py b/nemo_rl/distributed/model_utils.py index 18b387e82a4..9dc4553fef8 100644 --- a/nemo_rl/distributed/model_utils.py +++ b/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, diff --git a/nemo_rl/distributed/ray_actor_environment_registry.py b/nemo_rl/distributed/ray_actor_environment_registry.py index 3ddba6490e4..26489dc6c52 100644 --- a/nemo_rl/distributed/ray_actor_environment_registry.py +++ b/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, diff --git a/nemo_rl/models/automodel/data.py b/nemo_rl/models/automodel/data.py index 98eed48d4fe..acb9c51a8c8 100644 --- a/nemo_rl/models/automodel/data.py +++ b/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 diff --git a/nemo_rl/models/automodel/diffusion_attention.py b/nemo_rl/models/automodel/diffusion_attention.py new file mode 100644 index 00000000000..04e2622a913 --- /dev/null +++ b/nemo_rl/models/automodel/diffusion_attention.py @@ -0,0 +1,216 @@ +# 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. + +"""AutoModel attention helpers for completion-only diffusion replay. + +The replay layout is ``[noisy_response | clean_prompt_response]``. The published +Nemotron-Labs-Diffusion Hugging Face attention implements only the older symmetric +``[x_t | x_0]`` layout: it assumes equal halves, uses one scalar split point, and +does not mask per-sample padding. These helpers provide the asymmetric mask and +segmented position ids used by the Megatron-Bridge reference implementation. + +The worker installs the resulting ``BlockMask`` through the model's existing +``sbd_block_diff_mask`` cache. This is a narrowly checked compatibility adapter; +models that expose a native ``set_asymmetric_ar_metadata`` API use that instead. +""" + +from collections.abc import Callable +from typing import Any + +import torch +from torch.nn.attention.flex_attention import BlockMask, create_block_mask + + +def asymmetric_semi_ar_mask_mod( + *, + block_size: int, + noisy_length: int, + noisy_response_offset: int, + prompt_lengths: torch.Tensor, + noisy_valid_lengths: torch.Tensor, + clean_lengths: torch.Tensor, +) -> Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]: + """Return the batch-aware asymmetric semi-AR attention predicate.""" + + def asymmetric_semi_ar_mask( + batch_idx: torch.Tensor, + head_idx: torch.Tensor, + query_idx: torch.Tensor, + kv_idx: torch.Tensor, + ) -> torch.Tensor: + del head_idx + prompt_length = prompt_lengths[batch_idx] + noisy_valid_length = noisy_valid_lengths[batch_idx] + clean_length = clean_lengths[batch_idx] + + query_is_noisy = query_idx < noisy_length + kv_is_noisy = kv_idx < noisy_length + query_noisy_relative = query_idx - noisy_response_offset + kv_noisy_relative = kv_idx - noisy_response_offset + query_noisy_valid = ( + query_is_noisy + & (query_noisy_relative >= 0) + & (query_noisy_relative < noisy_valid_length) + ) + kv_noisy_valid = ( + kv_is_noisy + & (kv_noisy_relative >= 0) + & (kv_noisy_relative < noisy_valid_length) + ) + + query_block = torch.div(query_noisy_relative, block_size, rounding_mode="floor") + kv_block = torch.div(kv_noisy_relative, block_size, rounding_mode="floor") + noisy_same_block = ( + query_noisy_valid & kv_noisy_valid & (query_block == kv_block) + ) + + query_clean_idx = query_idx - noisy_length + kv_clean_idx = kv_idx - noisy_length + query_clean_valid = ( + (~query_is_noisy) + & (query_clean_idx >= 0) + & (query_clean_idx < clean_length) + ) + kv_clean_valid = ( + (~kv_is_noisy) & (kv_clean_idx >= 0) & (kv_clean_idx < clean_length) + ) + + clean_prompt = kv_clean_valid & (kv_clean_idx < prompt_length) + clean_response_relative = kv_clean_idx - prompt_length + clean_previous_response_blocks = ( + kv_clean_valid + & (clean_response_relative >= 0) + & (clean_response_relative < query_block * block_size) + ) + noisy_query = query_noisy_valid & ( + noisy_same_block | clean_prompt | clean_previous_response_blocks + ) + + clean_causal = ( + query_clean_valid & kv_clean_valid & (kv_clean_idx <= query_clean_idx) + ) + valid_query = query_noisy_valid | query_clean_valid + invalid_query_self = (~valid_query) & (query_idx == kv_idx) + return noisy_query | clean_causal | invalid_query_self + + return asymmetric_semi_ar_mask + + +def build_asymmetric_semi_ar_block_mask( + *, + block_size: int, + noisy_length: int, + clean_length: int, + noisy_response_offset: int, + prompt_lengths: torch.Tensor, + noisy_valid_lengths: torch.Tensor, + clean_lengths: torch.Tensor, +) -> BlockMask: + """Build the compact flex-attention mask for one processed microbatch.""" + _validate_metadata( + prompt_lengths=prompt_lengths, + noisy_valid_lengths=noisy_valid_lengths, + clean_lengths=clean_lengths, + ) + total_length = noisy_length + clean_length + mask_mod = asymmetric_semi_ar_mask_mod( + block_size=block_size, + noisy_length=noisy_length, + noisy_response_offset=noisy_response_offset, + prompt_lengths=prompt_lengths, + noisy_valid_lengths=noisy_valid_lengths, + clean_lengths=clean_lengths, + ) + return create_block_mask( + mask_mod, + B=prompt_lengths.shape[0], + H=None, + Q_LEN=total_length, + KV_LEN=total_length, + device=prompt_lengths.device, + ) + + +def build_asymmetric_position_ids( + *, + noisy_length: int, + clean_length: int, + noisy_response_offset: int, + prompt_lengths: torch.Tensor, + noisy_valid_lengths: torch.Tensor, +) -> torch.Tensor: + """Build per-sample RoPE positions for ``[noisy | clean]`` replay.""" + batch_size = prompt_lengths.shape[0] + device = prompt_lengths.device + noisy_positions = torch.arange(noisy_length, device=device).unsqueeze(0) + noisy_positions = noisy_positions.expand(batch_size, -1) + noisy_relative = noisy_positions - noisy_response_offset + noisy_valid = (noisy_relative >= 0) & ( + noisy_relative < noisy_valid_lengths.unsqueeze(1) + ) + noisy_position_ids = prompt_lengths.unsqueeze(1) + noisy_relative.clamp_min(0) + noisy_position_ids = torch.where( + noisy_valid, noisy_position_ids, torch.zeros_like(noisy_position_ids) + ) + clean_position_ids = torch.arange(clean_length, device=device).unsqueeze(0) + clean_position_ids = clean_position_ids.expand(batch_size, -1) + return torch.cat((noisy_position_ids, clean_position_ids), dim=1) + + +def install_hf_nld_asymmetric_mask( + module: Any, + *, + block_mask: BlockMask, +) -> bool: + """Install a mask into the published HF NLD flex-attention cache. + + Returns ``False`` for an unknown attention implementation so callers can fail + instead of silently falling back to its symmetric mask. + """ + if module.__class__.__name__ != "MinistralFlexAttention": + return False + required_attributes = ( + "sbd_block_diff_mask", + "block_size_orig", + "set_attention_mode", + ) + if not all(hasattr(module, name) for name in required_attributes): + return False + module.sbd_block_diff_mask = block_mask + return True + + +def clear_hf_nld_asymmetric_mask(module: Any) -> bool: + """Clear a mask installed by :func:`install_hf_nld_asymmetric_mask`.""" + if module.__class__.__name__ != "MinistralFlexAttention" or not hasattr( + module, "sbd_block_diff_mask" + ): + return False + module.sbd_block_diff_mask = None + return True + + +def _validate_metadata( + *, + prompt_lengths: torch.Tensor, + noisy_valid_lengths: torch.Tensor, + clean_lengths: torch.Tensor, +) -> None: + tensors = (prompt_lengths, noisy_valid_lengths, clean_lengths) + if any(tensor.ndim != 1 for tensor in tensors): + raise ValueError("Asymmetric semi-AR attention metadata must be 1D") + if not (prompt_lengths.shape == noisy_valid_lengths.shape == clean_lengths.shape): + raise ValueError( + "Asymmetric semi-AR attention metadata tensors must have matching shapes" + ) diff --git a/nemo_rl/models/automodel/diffusion_train.py b/nemo_rl/models/automodel/diffusion_train.py new file mode 100644 index 00000000000..d6068555301 --- /dev/null +++ b/nemo_rl/models/automodel/diffusion_train.py @@ -0,0 +1,236 @@ +# 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. + +"""Automodel (DTensor) post-processors for diffusion-LLM RL. + +This is the Automodel counterpart of ``nemo_rl/models/megatron/diffu_grpo_train.py``. +It supplies the two post-processors a diffusion policy worker needs, differing from +the autoregressive ones in exactly one respect: + + An AR model predicts token ``t+1`` at position ``t``, so its logprobs are read + with a one-position shift. A diffusion model predicts token ``t`` AT position + ``t``, so the target is read at the SAME position -- no shift. + +Both classes SUBCLASS the stock Automodel post-processors rather than replacing +them. That matters: ``forward_with_post_processing_fn`` dispatches with +``isinstance``, so a subclass matches the existing branch and needs no change to +the ``PostProcessingFunction`` union or the dispatch chain. +""" + +from typing import Any, Optional + +import torch + +from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.models.automodel.data import ProcessedInputs +from nemo_rl.models.automodel.train import ( + LogprobsPostProcessor, + LossPostProcessor, + apply_top_k_top_p_filtering_for_local_logits, +) + + +def same_position_logprobs( + logits: torch.Tensor, + target_ids: torch.Tensor, + *, + exclude_token_id: Optional[int] = None, + temperature: Optional[float] = None, + sampling_params: Optional[TrainingSamplingParams] = None, +) -> torch.Tensor: + """Log-probability of ``target_ids[b, t]`` under ``logits[b, t]`` -- no shift. + + The autoregressive counterpart rolls the targets left by one and drops a column. + Here position ``t`` scores the token that belongs at position ``t``, which is the + whole point of a masked-diffusion objective. + + Args: + logits: ``[B, S, V]`` model output. + target_ids: ``[B, S]`` tokens to score, aligned to ``logits``. + exclude_token_id: If given, that vocabulary entry is removed from the + softmax before normalising. Used to drop the MASK token, which the + SGLang FastDiffuser decoder also excludes from its x0 distribution -- + train and generate must agree or generation-KL blows up. + temperature: If given, logits are divided by it first. Must match whatever + was applied when ``prev_logprobs`` were computed. + sampling_params: Optional top-k/top-p filtering parameters. Temperature in + this object is intentionally ignored; callers scale exactly once before + invoking this helper. + + Returns: + ``[B, S]`` log probabilities. + """ + if logits.ndim != 3: + raise ValueError(f"logits must be [B, S, V], got {tuple(logits.shape)}") + if target_ids.shape != logits.shape[:2]: + raise ValueError( + f"target_ids {tuple(target_ids.shape)} must match logits prefix " + f"{tuple(logits.shape[:2])}" + ) + + if temperature is not None: + logits = logits / temperature + + if exclude_token_id is not None: + # Mask in a copy; -inf drops the entry from the softmax denominator. + logits = logits.clone() + logits[..., exclude_token_id] = float("-inf") + + logits = apply_top_k_top_p_filtering_for_local_logits(logits, sampling_params) + + logprobs = torch.log_softmax(logits.to(torch.float32), dim=-1) + return logprobs.gather(dim=-1, index=target_ids.unsqueeze(-1)).squeeze(-1) + + +class DiffusionLogprobsPostProcessor(LogprobsPostProcessor): + """Same-position logprobs for a diffusion objective. + + Reads the tokens to score from ``data_dict['diffu_grpo_target_ids']`` when + present -- a diffusion batch feeds MASKED tokens to the model but scores the + CLEAN tokens underneath, so input and target genuinely differ. Falls back to + ``input_ids`` otherwise. + """ + + def __init__( + self, *args: Any, exclude_token_id: Optional[int] = None, **kwargs: Any + ): + super().__init__(*args, **kwargs) + self.exclude_token_id = exclude_token_id + + def __call__( + self, + logits: torch.Tensor, + data_dict: BatchedDataDict[Any], + processed_inputs: ProcessedInputs, + original_batch_size: int, + original_seq_len: int, + sequence_dim: int = 1, + ) -> torch.Tensor: + if self.cp_size > 1: + raise NotImplementedError( + "Context parallelism is not supported for the diffusion logprob " + "path. The default Automodel CP route requires is_causal=True and " + "nulls the attention mask, which a block-diffusion objective cannot " + "use. See the DSV4 manual-CP precedent for how a model-specific CP " + "path would be added." + ) + + target_ids = data_dict.get("diffu_grpo_target_ids", None) + if target_ids is None: + target_ids = processed_inputs.input_ids + + # ``forward_with_post_processing_fn`` applies temperature scaling before + # dispatching every LogprobsPostProcessor subclass. Applying it again here + # would score at temperature squared when T != 1. + return same_position_logprobs( + logits, + target_ids, + exclude_token_id=self.exclude_token_id, + sampling_params=self.sampling_params, + ) + + +class DiffusionLossPostProcessor(LossPostProcessor): + """Clipped-PG loss over position-aligned (unshifted) diffusion logprobs. + + The stock LossPostProcessor hands raw logits to ``ClippedPGLossFn.__call__``, + which internally slices five tensors by ``[:, 1:]`` to meet the next-token + contract. That is wrong here: a diffusion model scores token ``t`` AT position + ``t``, so nothing should be shifted. + + We therefore compute the logprobs ourselves and call the loss's + ``compute_from_aligned_tensors`` entry point, which takes already-aligned + tensors. That method lives on the SHARED loss class, so the Megatron and + Automodel diffusion paths use exactly the same objective code. + + Note the denominator: ``global_valid_toks`` as computed by + ``process_global_batch`` counts ``token_mask[:, 1:]`` -- an AR next-token + count. A diffusion objective supervises a different set of positions, so + callers pass the harvested-token count through ``metadata`` instead. If it is + absent we fall back to the AR count and warn, rather than silently scaling the + gradient wrong. + """ + + def __init__( + self, + *args: Any, + exclude_token_id: Optional[int] = None, + valid_toks_override: Optional[torch.Tensor] = None, + **kwargs: Any, + ): + super().__init__(*args, **kwargs) + self.exclude_token_id = exclude_token_id + self.valid_toks_override = valid_toks_override + + def __call__( + self, + logits: torch.Tensor, + data_dict: BatchedDataDict[Any], + processed_inputs: ProcessedInputs, + global_valid_seqs: torch.Tensor, + global_valid_toks: torch.Tensor, + sequence_dim: int = 1, + ) -> tuple[torch.Tensor, dict[str, Any]]: + if self.cp_size > 1: + raise NotImplementedError( + "Context parallelism is not supported for the diffusion loss path." + ) + + target_ids = data_dict.get("diffu_grpo_target_ids", None) + if target_ids is None: + target_ids = processed_inputs.input_ids + + # Temperature was already applied by ``forward_with_post_processing_fn``. + curr_logprobs = same_position_logprobs( + logits, + target_ids, + exclude_token_id=self.exclude_token_id, + sampling_params=self.sampling_params, + ) + + # Only the NOISY half is scored. The clean half exists solely to supply + # previous-block context, so every per-token tensor is truncated to it -- + # mirroring nemo_rl/models/megatron/diffu_grpo_train.py so both backends + # feed the shared loss identical inputs. + noisy_length = int(data_dict["diffu_grpo_noisy_lengths"][0].item()) + + def _clip(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + return None if t is None else t[:, :noisy_length] + + curr_logprobs = curr_logprobs[:, :noisy_length] + + # Note this is diffu_grpo_loss_mask, NOT the generic token_mask: the set of + # supervised positions is determined by the reveal schedule, not by which + # tokens are non-padding. + loss_mask = _clip(data_dict["diffu_grpo_loss_mask"]) + + valid_toks = ( + self.valid_toks_override + if self.valid_toks_override is not None + else global_valid_toks + ) + + loss, metrics = self.loss_fn.compute_from_aligned_tensors( + curr_logprobs=curr_logprobs, + token_mask=loss_mask, + sample_mask=data_dict["sample_mask"], + advantages=_clip(data_dict["advantages"]), + prev_logprobs=_clip(data_dict["prev_logprobs"]), + generation_logprobs=_clip(data_dict["generation_logprobs"]), + reference_policy_logprobs=_clip(data_dict.get("reference_policy_logprobs")), + global_valid_seqs=global_valid_seqs, + global_valid_toks=valid_toks, + ) + return loss, metrics diff --git a/nemo_rl/models/automodel/setup.py b/nemo_rl/models/automodel/setup.py index a539560b0d4..af6e99cd6d8 100644 --- a/nemo_rl/models/automodel/setup.py +++ b/nemo_rl/models/automodel/setup.py @@ -441,6 +441,13 @@ def setup_distributed( param_dtype=dtype, reduce_dtype=torch.float32, output_dtype=torch.float32, + # Keep inputs to recursively sharded transformer blocks aligned + # with their compute parameters. In particular, HF rotary + # embeddings may be produced by an unsharded FP32 root module; + # leaving them in FP32 promotes BF16 Q/K while V remains BF16, + # which fused attention rejects. This also matches AutoModel's + # default FSDP2 mixed-precision policy. + cast_forward_inputs=True, ), offload_policy=CPUOffloadPolicy(pin_memory=False) if cpu_offload else None, activation_checkpointing=config["dtensor_cfg"]["activation_checkpointing"], diff --git a/nemo_rl/models/automodel/train.py b/nemo_rl/models/automodel/train.py index d2b59794007..0ad20407096 100644 --- a/nemo_rl/models/automodel/train.py +++ b/nemo_rl/models/automodel/train.py @@ -86,6 +86,14 @@ def model_forward( use_cache=False, ) + reserved_model_kwargs = processed_inputs.model_kwargs.keys() & model_args.keys() + if reserved_model_kwargs: + raise ValueError( + "ProcessedInputs.model_kwargs cannot override core model arguments: " + f"{sorted(reserved_model_kwargs)}" + ) + model_args.update(processed_inputs.model_kwargs) + # Add flash attention kwargs if applicable if processed_inputs.has_flash_attention: model_args["flash_attn_kwargs"] = processed_inputs.flash_attn_kwargs diff --git a/nemo_rl/models/policy/workers/block_just_grpo_dtensor_policy_worker.py b/nemo_rl/models/policy/workers/block_just_grpo_dtensor_policy_worker.py new file mode 100644 index 00000000000..0210cf9202f --- /dev/null +++ b/nemo_rl/models/policy/workers/block_just_grpo_dtensor_policy_worker.py @@ -0,0 +1,292 @@ +# 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. + +"""BlockJustGRPO(-Fast) on the Automodel (DTensor) backend. + +Automodel counterpart of ``block_just_grpo_megatron_policy_worker``. The estimator +math is imported unchanged from ``nemo_rl.algorithms.block_just_grpo_logprobs``, +which has no backend imports -- only the driving of forwards differs between the two +backends, which is the whole point of the hook set. + +Scoring cost, for orientation (block_size=16, 512-token response, 32 blocks): + + naive 512 forwards one per token + Block 16 forwards score offset j of EVERY block at once (exact) + Fast 4 forwards keep the top ceil(0.25*16) highest-entropy offsets + +"Fast" needs per-token entropy from the generation engine. Its absence is a +fail-fast (``require_generation_entropy``) rather than a silent zero-fill, because +sparsifying on zeros would rank positions arbitrarily and still train. +""" + +from typing import Any, Optional + +import ray +import torch + +from nemo_rl.algorithms.block_just_grpo_logprobs import ( + BlockJustGRPORevealSchedule, + build_block_kept_mask, + build_block_reveal_base, + get_block_reveal_logprob_estimation_cfg, + make_reveal_level_view, + scatter_block_reveal_logprobs, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.models.policy import PolicyConfig +from nemo_rl.models.policy.interfaces import LogprobOutputSpec +from nemo_rl.models.policy.utils import get_runtime_env_for_policy_worker +from nemo_rl.models.policy.workers.diffusion_dtensor_policy_worker import ( + DiffusionDTensorPolicyWorkerImpl, +) + + +class BlockJustGRPODTensorPolicyWorkerImpl(DiffusionDTensorPolicyWorkerImpl): + """Leftmost-reveal logprobs computed in ``block_size`` (or top-m) passes.""" + + # ---- config ------------------------------------------------------------ + + def _validate_diffusion_algorithm_support(self) -> None: + self._validate_diffusion_support("BlockJustGRPO") + paradigm = (self.cfg.get("hf_config_overrides") or {}).get( + "dlm_paradigm" + ) or getattr(getattr(self, "model_config", None), "dlm_paradigm", None) + if paradigm != "sbd_block_diff": + raise NotImplementedError( + "BlockJustGRPO completion-only replay on AutoModel requires " + "dlm_paradigm='sbd_block_diff'; the block_diff cache follows a " + "different symmetric attention contract." + ) + get_block_reveal_logprob_estimation_cfg(self.cfg) + + def _block_reveal_cfg(self) -> Any: + return get_block_reveal_logprob_estimation_cfg(self.cfg) + + def _block_reveal_block_size(self) -> int: + cfg = self._block_reveal_cfg() + return int(cfg.get("block_size") or self._diffusion_block_size()) + + def _block_reveal_tokens_per_level(self) -> int: + return max(1, int(self._block_reveal_cfg().get("reveal_tokens_per_level") or 1)) + + def _mask_token_id(self) -> int: + return int(self._block_reveal_cfg()["mask_token_id"]) + + def _exclude_mask_token(self) -> bool: + return bool( + self._block_reveal_cfg().get("exclude_mask_token_from_logits", False) + ) + + # ---- post-processors --------------------------------------------------- + + def _make_logprobs_post_processor(self, **kwargs: Any) -> Any: + from nemo_rl.models.automodel.diffusion_train import ( + DiffusionLogprobsPostProcessor, + ) + + return DiffusionLogprobsPostProcessor( + exclude_token_id=self._mask_token_id() + if self._exclude_mask_token() + else None, + **kwargs, + ) + + def _make_loss_post_processor(self, **kwargs: Any) -> Any: + from nemo_rl.models.automodel.diffusion_train import DiffusionLossPostProcessor + + return DiffusionLossPostProcessor( + exclude_token_id=self._mask_token_id() + if self._exclude_mask_token() + else None, + valid_toks_override=getattr(self, "_fast_valid_toks", None), + **kwargs, + ) + + # ---- training ---------------------------------------------------------- + + def _build_training_batch( + self, data: BatchedDataDict[Any], mbs: int + ) -> tuple[BatchedDataDict[Any], PolicyConfig, int, dict[str, Any]]: + cfg = self._block_reveal_cfg() + block_size = self._block_reveal_block_size() + reveal_k = self._block_reveal_tokens_per_level() + + base, _n_samples, num_levels, selected_offsets = build_block_reveal_base( + data, + mask_token_id=cfg["mask_token_id"], + pad_token_id=self.tokenizer.pad_token_id, + block_size=block_size, + pad_to_length=None, + include_loss=True, + max_reveal_levels=cfg.get("max_reveal_levels"), + reveal_tokens_per_level=reveal_k, + fast_entropy_level_ratio=cfg.get("fast_entropy_level_ratio"), + eos_token_id=self.tokenizer.eos_token_id, + force_eos=cfg.get("fast_force_eos", True), + ) + + # A schedule, not a batch. It overrides only .size and + # .make_microbatch_iterator -- exactly the two members the Automodel + # microbatch iterator uses -- so one logical batch expands into + # num_levels x microbatches forwards, all accumulating into ONE optimizer + # step. This is where Automodel is easier than Megatron: its + # forward/backward is a plain Python loop rather than an mcore schedule. + schedule = BlockJustGRPORevealSchedule(base).configure( + num_levels=num_levels, + block_size=block_size, + harvest_keys=("diffu_grpo_score_mask", "diffu_grpo_loss_mask"), + reveal_tokens_per_level=reveal_k, + selected_offsets=selected_offsets, + ) + + metadata: dict[str, Any] = {"num_levels": num_levels} + if selected_offsets is not None: + # Fast harvests only a subset of positions, so the loss must be + # normalised by the KEPT count. Using the full response length would + # silently scale the gradient with fast_entropy_level_ratio. + metadata["fast_global_valid_toks"] = self._fast_global_valid_toks( + base, selected_offsets, block_size + ) + self._fast_valid_toks = metadata["fast_global_valid_toks"] + return schedule, self.cfg, mbs, metadata + + def _fast_global_valid_toks( + self, + base: BatchedDataDict[Any], + selected_offsets: torch.Tensor, + block_size: int, + ) -> torch.Tensor: + """DP-reduced count of positions Fast actually harvests. + + The Megatron worker's single mcore call + (``parallel_state.get_data_parallel_group()``) becomes ``self.dp_mesh`` + here -- the only backend-specific line in this class. + """ + local_kept = build_block_kept_mask(base, selected_offsets, block_size).sum() + to_reduce = local_kept.detach().to(dtype=torch.float32).reshape(1).cuda() + torch.distributed.all_reduce(to_reduce, group=self._dp_group()) + return to_reduce[0] + + # ---- logprobs ---------------------------------------------------------- + + def get_logprobs( + self, data: BatchedDataDict[Any], micro_batch_size: Optional[int] = None + ) -> BatchedDataDict[LogprobOutputSpec]: + """One forward per reveal level; per-level logprobs are summed. + + Only ONE reveal level is resident at a time, which is what keeps long + contexts tractable. + """ + self._validate_diffusion_algorithm_support() + cfg = self._block_reveal_cfg() + block_size = self._block_reveal_block_size() + reveal_k = self._block_reveal_tokens_per_level() + reveal_mbs = micro_batch_size or self.cfg["logprob_batch_size"] + + base, num_samples, num_levels, selected_offsets = build_block_reveal_base( + data, + mask_token_id=cfg["mask_token_id"], + pad_token_id=self.tokenizer.pad_token_id, + block_size=block_size, + pad_to_length=None, + include_loss=False, + max_reveal_levels=cfg.get("max_reveal_levels"), + reveal_tokens_per_level=reveal_k, + fast_entropy_level_ratio=cfg.get("fast_entropy_level_ratio"), + eos_token_id=self.tokenizer.eos_token_id, + force_eos=cfg.get("fast_force_eos", True), + ) + + if num_levels == 0: + empty = torch.zeros_like(data["input_ids"], dtype=torch.float32) + return BatchedDataDict[LogprobOutputSpec](logprobs=empty).to("cpu") + + # Per-level state. NOTE: this makes get_logprobs non-reentrant -- safe + # only with one in-flight call per worker. The reference-logprob pass + # re-enters this method under a weight swap, but sequentially. + self._br_base = base + self._br_selected_offsets = selected_offsets + self._br_block_size_cur = block_size + self._br_reveal_k = reveal_k + self._br_reveal_mbs = reveal_mbs + self._br_num_samples = num_samples + self._br_original_seq_len = int(data["input_ids"].shape[1]) + + accumulated: Optional[torch.Tensor] = None + try: + for level in range(num_levels): + self._br_level = level + level_out = super().get_logprobs( + data=data, micro_batch_size=reveal_mbs + )["logprobs"] + accumulated = ( + level_out if accumulated is None else accumulated + level_out + ) + finally: + self._br_base = None + + return BatchedDataDict[LogprobOutputSpec](logprobs=accumulated).to("cpu") + + def _build_logprob_batch( + self, data: BatchedDataDict[Any], micro_batch_size: Optional[int] + ) -> tuple[Optional[BatchedDataDict[Any]], PolicyConfig, int, dict[str, Any]]: + """Build the view for the reveal level the outer loop is currently on.""" + view = make_reveal_level_view( + self._br_base, + self._br_level, + self._br_block_size_cur, + ("diffu_grpo_score_mask",), + self._br_reveal_k, + selected_offsets=self._br_selected_offsets, + ) + return ( + view, + self.cfg, + micro_batch_size or self._br_reveal_mbs, + { + "num_samples": self._br_num_samples, + "original_seq_len": self._br_original_seq_len, + "noisy_response_offset": int( + view["diffu_grpo_noisy_response_offsets"][0].item() + ), + }, + ) + + def _finalize_logprobs_from_outputs( + self, + list_of_logprobs: list[dict[str, torch.Tensor]], + *, + original_data: BatchedDataDict[Any], + transformed_data: BatchedDataDict[Any], + metadata: dict[str, Any], + ) -> torch.Tensor: + flat = torch.cat([lp["logprobs"] for lp in list_of_logprobs], dim=0) + return scatter_block_reveal_logprobs( + flat_logprobs=flat, + harvest_mask=transformed_data["block_reveal_harvest_mask"], + sample_index=transformed_data["block_reveal_sample_index"], + completion_starts=transformed_data["diffu_grpo_completion_starts"], + noisy_response_offset=metadata["noisy_response_offset"], + original_seq_len=metadata["original_seq_len"], + num_samples=metadata["num_samples"], + ) + + +@ray.remote( + runtime_env=get_runtime_env_for_policy_worker( + "block_just_grpo_dtensor_policy_worker" + ) +) # pragma: no cover +class BlockJustGRPODTensorPolicyWorker(BlockJustGRPODTensorPolicyWorkerImpl): + pass diff --git a/nemo_rl/models/policy/workers/diffusion_dtensor_policy_worker.py b/nemo_rl/models/policy/workers/diffusion_dtensor_policy_worker.py new file mode 100644 index 00000000000..2cd09971aea --- /dev/null +++ b/nemo_rl/models/policy/workers/diffusion_dtensor_policy_worker.py @@ -0,0 +1,647 @@ +# 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. + +"""Automodel (DTensor) base worker for diffusion-LLM RL. + +Counterpart of ``diffusion_megatron_policy_worker.DiffusionMegatronPolicyWorkerImpl``, +re-parented onto ``DTensorPolicyWorkerV2Impl``. + +The hook set is deliberately IDENTICAL in shape to the Megatron ABC's, because that +hook set is already validated: seven estimators (JustGRPO, DiffuGRPO, BlockJustGRPO, +CoupledGRPO, ESPO, TraceGRPO, hybrid AR+diffusion) sit on it without the base class +knowing about any of them. Names drop the ``megatron`` infix since they are no +longer backend-specific. + +What differs from the Megatron ABC, and why: + +* No pipeline-parallel handling. Automodel hardcodes ``pp_size=1``, so the + last-stage broadcasts and cross-stage plumbing simply do not exist here. Net + deletion. +* No ``rerun_state_machine``, no ``StragglerDetector``, no ``zero_grad_buffer``. +* ``stream_weights_via_http`` is NOT overridden. The Megatron ABC overrides it to + call ``_iter_params_with_optional_kv_scales()``, which exists only on the + Megatron worker -- inheriting that override here would AttributeError on the + first refit. The DTensor base already implements refit correctly. +* Attention control is a hook rather than a duck-typed walk over ``model.modules()``. + +Subclasses implement the six abstract hooks. Everything else is inherited. +""" + +from abc import ABC, abstractmethod +from contextlib import AbstractContextManager, contextmanager, nullcontext +from typing import Any, Iterator, Optional +import warnings + +import torch +from nemo_automodel.components.training.utils import scale_grads_and_clip_grad_norm + +from nemo_rl.algorithms.loss.interfaces import LossFunction +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.models.automodel.data import ( + ProcessedMicrobatch, + check_sequence_dim, + get_microbatch_iterator, + process_global_batch, +) +from nemo_rl.models.automodel.diffusion_attention import ( + build_asymmetric_position_ids, + build_asymmetric_semi_ar_block_mask, + clear_hf_nld_asymmetric_mask, + install_hf_nld_asymmetric_mask, +) +from nemo_rl.models.automodel.train import ( + aggregate_training_statistics, + automodel_forward_backward, + forward_with_post_processing_fn, +) +from nemo_rl.models.policy import PolicyConfig +from nemo_rl.models.policy.interfaces import LogprobOutputSpec +from nemo_rl.models.policy.workers.dtensor_policy_worker_v2 import ( + DTensorPolicyWorkerV2Impl, + get_train_context, +) +from nemo_rl.utils.nsys import wrap_with_nvtx_name + + +class DiffusionDTensorPolicyWorkerImpl(DTensorPolicyWorkerV2Impl, ABC): + """Shared machinery for masked-diffusion RL estimators on the Automodel backend.""" + + # ---- capability guards ------------------------------------------------ + + def _validate_diffusion_support(self, algorithm_name: str) -> None: + """Reject configurations this path cannot honour, loudly and early.""" + dtensor_cfg = self.cfg.get("dtensor_cfg", {}) + + if dtensor_cfg.get("context_parallel_size", 1) != 1: + raise NotImplementedError( + f"{algorithm_name}: context_parallel_size > 1 is not supported on the " + "Automodel backend. The default CP route requires is_causal=True and " + "nulls the attention mask (automodel/data.py), which a block-diffusion " + "objective cannot use. A model-specific manual-CP path would be needed " + "-- see NRL_DSV4_MANUAL_CP for the precedent." + ) + + if dtensor_cfg.get("tensor_parallel_size", 1) != 1: + raise NotImplementedError( + f"{algorithm_name}: tensor_parallel_size > 1 is not supported yet. " + "The diffusion post-processors currently require a full local " + "vocabulary for same-position logprobs." + ) + + if self.cfg.get("dynamic_batching", {}).get("enabled", False): + raise NotImplementedError( + f"{algorithm_name}: dynamic batching is not supported. Reveal-level " + "schedules currently implement the fixed-size microbatch iterator only." + ) + + self._validate_attention_paradigm() + + seq_packing = self.cfg.get("sequence_packing", {}) + if seq_packing.get("enabled", False): + raise NotImplementedError( + f"{algorithm_name}: sequence packing is not supported. Packing keeps " + "documents apart via cu_seqlens precisely BECAUSE attention is causal; " + "bidirectional attention inside a packed buffer leaks across document " + "boundaries unless the mask is rebuilt block-diagonally." + ) + + @abstractmethod + def _validate_diffusion_algorithm_support(self) -> None: + """Estimator-specific config validation.""" + + # ---- attention control ------------------------------------------------ + # + # The Megatron ABC discovers attention modules by duck-typing over + # model.modules() for set_inference_mode/set_inference_params/clear_kv_cache. + # The HF modeling code exposes set_attention_mode(mode, block_size) on its + # flex-attention class, so we drive that directly instead. + + #: dlm_paradigm values that cause the model to build flex attention with the + #: asymmetric [noisy | clean] mask. Anything else gives plain attention. + FLEX_PARADIGMS = ("block_diff", "sbd_block_diff") + + def _validate_attention_paradigm(self) -> None: + """Fail if the model was built with plain attention. + + VERIFIED ON HARDWARE: MinistralDiffEncoderModel.__init__ selects its + attention class from ``dlm_paradigm``. The shipped checkpoint declares + ``autoregressive``, which yields plain Ministral3Attention with + ``diffusion_lm=False`` -- i.e. ordinary causal attention, with NO error. + Training would converge on the wrong objective. + + Selection therefore happens at config time: + + policy: + hf_config_overrides: + dlm_paradigm: sbd_block_diff + + This check exists because that failure is otherwise silent. + """ + overrides = self.cfg.get("hf_config_overrides") or {} + paradigm = overrides.get("dlm_paradigm") or getattr( + getattr(self, "model_config", None), "dlm_paradigm", None + ) + if paradigm not in self.FLEX_PARADIGMS: + raise ValueError( + f"dlm_paradigm={paradigm!r} builds plain (causal) attention. A " + f"masked-diffusion objective needs one of {self.FLEX_PARADIGMS}. " + "Set policy.hf_config_overrides.dlm_paradigm=sbd_block_diff. " + "Without this the run trains successfully on the wrong objective." + ) + if not self._diffusion_attention_modules(): + raise RuntimeError( + f"dlm_paradigm={paradigm!r} was requested but no attention module " + "exposes set_attention_mode(); flex attention was not constructed." + ) + + def _diffusion_attention_modules(self) -> list[Any]: + return [m for m in self.model.modules() if hasattr(m, "set_attention_mode")] + + @contextmanager + def _attention_mode(self, mode: Optional[str], block_size: Optional[int] = None): + """Temporarily switch every attention module to ``mode``.""" + if mode is None: + yield + return + modules = self._diffusion_attention_modules() + if not modules: + raise RuntimeError( + "No attention module exposes set_attention_mode(). The diffusion " + "path needs to control the attention regime; the loaded model does " + "not appear to be a diffusion LM." + ) + previous = [ + (m, getattr(m, "mode", None), getattr(m, "block_size", None)) + for m in modules + ] + try: + for m in modules: + m.set_attention_mode(mode, block_size=block_size) + yield + finally: + for m, prev_mode, prev_bs in previous: + if prev_mode is not None: + m.set_attention_mode(prev_mode, block_size=prev_bs) + + def _training_attention_context(self) -> AbstractContextManager[Any]: + overrides = self.cfg.get("hf_config_overrides") or {} + return self._attention_mode( + overrides.get("dlm_paradigm"), block_size=self._diffusion_block_size() + ) + + def _logprob_attention_context(self) -> AbstractContextManager[Any]: + overrides = self.cfg.get("hf_config_overrides") or {} + return self._attention_mode( + overrides.get("dlm_paradigm"), block_size=self._diffusion_block_size() + ) + + # ---- execution --------------------------------------------------------- + + @wrap_with_nvtx_name("diffusion_dtensor_policy_worker/train") + def train( + self, + data: BatchedDataDict[Any], + loss_fn: LossFunction, + eval_mode: bool = False, + gbs: Optional[int] = None, + mbs: Optional[int] = None, + ) -> dict[str, Any]: + """Train over an estimator-provided batch or reveal-level schedule.""" + self._validate_diffusion_algorithm_support() + if gbs is None: + gbs = self.cfg["train_global_batch_size"] + if mbs is None: + mbs = self.cfg["train_micro_batch_size"] + + local_gbs = gbs // self.dp_size + total_dataset_size = torch.tensor(data.size, device="cuda") + torch.distributed.all_reduce( + total_dataset_size, + op=torch.distributed.ReduceOp.SUM, + group=self._dp_group(), + ) + num_global_batches = int(total_dataset_size.item()) // gbs + + if eval_mode: + ctx: AbstractContextManager[Any] = torch.no_grad() + self.model.eval() + else: + ctx = nullcontext() + self.model.train() + + empty_cache_steps = self.cfg.get("dtensor_cfg", {}).get( + "clear_cache_every_n_steps" + ) + if empty_cache_steps: + warnings.warn( + f"Emptying cache every {empty_cache_steps} microbatches; doing so " + "unnecessarily incurs a large performance overhead.", + stacklevel=2, + ) + + def on_microbatch_start(microbatch_idx: int) -> None: + if empty_cache_steps and microbatch_idx % empty_cache_steps == 0: + torch.cuda.empty_cache() + + data = data.to("cuda") + losses: list[float] = [] + all_mb_metrics: list[dict[str, Any]] = [] + grad_norm: Optional[torch.Tensor] = None + + with ctx, self._training_attention_context(): + for global_batch_idx in range(num_global_batches): + global_batch = process_global_batch( + data, + loss_fn, + self._dp_group(), + batch_idx=global_batch_idx, + batch_size=local_gbs, + ) + ( + transformed_data, + cfg_for_training, + train_mbs, + metadata, + ) = self._build_training_batch(global_batch["batch"], mbs) + sequence_dim, _ = check_sequence_dim(transformed_data) + + num_levels = int(metadata.get("num_levels", 1)) + samples_per_level = transformed_data.size // max(1, num_levels) + if samples_per_level % train_mbs != 0: + raise ValueError( + "The per-rank sample count for each diffusion reveal level " + f"({samples_per_level}) must be divisible by the training " + f"microbatch size ({train_mbs})." + ) + + global_valid_seqs = global_batch["global_valid_seqs"] + global_valid_toks = metadata.get( + "fast_global_valid_toks", global_batch["global_valid_toks"] + ) + + processed_iterator, iterator_len = get_microbatch_iterator( + transformed_data, + cfg_for_training, + train_mbs, + self.dp_mesh, + tokenizer=self.tokenizer, + cp_size=self.cp_size, + ) + processed_iterator = self._wrap_training_microbatch_iterator( + processed_iterator, metadata + ) + loss_post_processor = self._make_loss_post_processor( + loss_fn=loss_fn, + cfg=cfg_for_training, + device_mesh=self.device_mesh, + cp_mesh=self.cp_mesh, + tp_mesh=self.tp_mesh, + cp_size=self.cp_size, + dp_size=self.dp_size, + enable_seq_packing=self.enable_seq_packing, + sampling_params=self.sampling_params, + ) + + def train_context_fn(processed_inputs: Any): + return get_train_context( + cp_size=self.cp_size, + cp_mesh=self.cp_mesh, + cp_buffers=processed_inputs.cp_buffers, + sequence_dim=sequence_dim, + dtype=self.dtype, + autocast_enabled=self.autocast_enabled, + ) + + self.optimizer.zero_grad() + microbatch_results = automodel_forward_backward( + model=self.model, + data_iterator=processed_iterator, + post_processing_fn=loss_post_processor, + forward_only=eval_mode, + is_reward_model=False, + allow_flash_attn_args=self.allow_flash_attn_args, + global_valid_seqs=global_valid_seqs, + global_valid_toks=global_valid_toks, + sampling_params=self.sampling_params, + sequence_dim=sequence_dim, + dp_size=self.dp_size, + cp_size=self.cp_size, + num_global_batches=num_global_batches, + train_context_fn=train_context_fn, + num_valid_microbatches=iterator_len, + on_microbatch_start=on_microbatch_start, + ) + + microbatch_losses = [] + for microbatch_idx, (loss, loss_metrics) in enumerate( + microbatch_results + ): + if microbatch_idx >= iterator_len: + continue + loss_metrics["lr"] = self.optimizer.param_groups[0]["lr"] + loss_metrics["global_valid_seqs"] = global_valid_seqs.item() + loss_metrics["global_valid_toks"] = global_valid_toks.item() + if loss_metrics["num_valid_samples"] > 0: + microbatch_losses.append(loss.item()) + all_mb_metrics.append(loss_metrics) + + if not eval_mode: + raw_grad_norm = scale_grads_and_clip_grad_norm( + self.max_grad_norm, + [self.model], + norm_type=2.0, + pp_enabled=False, + device_mesh=self.device_mesh, + moe_mesh=self.moe_mesh, + ep_axis_name=( + "ep" + if self.moe_mesh is not None + and "ep" in self.moe_mesh.mesh_dim_names + else None + ), + pp_axis_name=None, + foreach=True, + num_label_tokens=1, + dp_group_size=self.dp_size * self.cp_size, + ) + grad_norm = torch.tensor( + raw_grad_norm, device="cpu", dtype=torch.float32 + ) + self.optimizer.step() + + losses.append(torch.tensor(microbatch_losses).sum().item()) + + self.optimizer.zero_grad() + if not eval_mode: + self.scheduler.step() + torch.cuda.empty_cache() + return aggregate_training_statistics( + losses=losses, + all_mb_metrics=all_mb_metrics, + grad_norm=grad_norm, + dp_group=self._dp_group(), + dtype=self.dtype, + ) + + @wrap_with_nvtx_name("diffusion_dtensor_policy_worker/get_logprobs") + def get_logprobs( + self, + data: BatchedDataDict[Any], + micro_batch_size: Optional[int] = None, + ) -> BatchedDataDict[LogprobOutputSpec]: + """Run estimator-specific diffusion forwards and restore ``[B, S]``.""" + self._validate_diffusion_algorithm_support() + ( + transformed_data, + cfg_for_logprobs, + logprob_mbs, + metadata, + ) = self._build_logprob_batch(data, micro_batch_size) + if transformed_data is None: + return BatchedDataDict[LogprobOutputSpec]( + logprobs=metadata["empty_logprobs"] + ).to("cpu") + + sequence_dim, _ = check_sequence_dim(transformed_data) + transformed_data = transformed_data.to("cuda") + self.model.eval() + processed_iterator, iterator_len = get_microbatch_iterator( + transformed_data, + cfg_for_logprobs, + logprob_mbs, + self.dp_mesh, + tokenizer=self.tokenizer, + cp_size=self.cp_size, + ) + processed_iterator = self._wrap_logprob_microbatch_iterator( + processed_iterator, metadata + ) + logprobs_post_processor = self._make_logprobs_post_processor( + cfg=cfg_for_logprobs, + device_mesh=self.device_mesh, + cp_mesh=self.cp_mesh, + tp_mesh=self.tp_mesh, + cp_size=self.cp_size, + enable_seq_packing=self.enable_seq_packing, + sampling_params=self.sampling_params, + ) + + list_of_logprobs: list[dict[str, torch.Tensor]] = [] + with torch.no_grad(), self._logprob_attention_context(): + for microbatch_idx, processed_microbatch in enumerate(processed_iterator): + processed_inputs = processed_microbatch.processed_inputs + with get_train_context( + cp_size=self.cp_size, + cp_mesh=self.cp_mesh, + cp_buffers=processed_inputs.cp_buffers, + sequence_dim=sequence_dim, + dtype=self.dtype, + autocast_enabled=self.autocast_enabled, + ): + token_logprobs, _metrics, _ = forward_with_post_processing_fn( + model=self.model, + post_processing_fn=logprobs_post_processor, + processed_mb=processed_microbatch, + is_reward_model=False, + allow_flash_attn_args=self.allow_flash_attn_args, + sampling_params=self.sampling_params, + sequence_dim=sequence_dim, + ) + if microbatch_idx < iterator_len: + list_of_logprobs.append({"logprobs": token_logprobs}) + + logprobs = self._finalize_logprobs_from_outputs( + list_of_logprobs, + original_data=data, + transformed_data=transformed_data, + metadata=metadata, + ) + return BatchedDataDict[LogprobOutputSpec](logprobs=logprobs).to("cpu") + + # ---- batch construction ------------------------------------------------ + + @abstractmethod + def _build_training_batch( + self, data: BatchedDataDict[Any], mbs: int + ) -> tuple[BatchedDataDict[Any], PolicyConfig, int, dict[str, Any]]: + """Return ``(batch_or_schedule, cfg, microbatch_size, metadata)``. + + Returning a SCHEDULE rather than a plain batch is what lets one logical + batch expand into N reveal-level forward passes: the schedule overrides + only ``.size`` and ``.make_microbatch_iterator``, which are exactly the two + members the Automodel microbatch iterator uses. + """ + + @abstractmethod + def _build_logprob_batch( + self, data: BatchedDataDict[Any], micro_batch_size: Optional[int] + ) -> tuple[Optional[BatchedDataDict[Any]], PolicyConfig, int, dict[str, Any]]: + """Same, for the logprob path.""" + + # ---- post-processors --------------------------------------------------- + + @abstractmethod + def _make_logprobs_post_processor(self, **kwargs: Any) -> Any: + """Return the LogprobsPostProcessor subclass this estimator needs.""" + + @abstractmethod + def _make_loss_post_processor(self, **kwargs: Any) -> Any: + """Return the LossPostProcessor subclass this estimator needs.""" + + @abstractmethod + def _finalize_logprobs_from_outputs( + self, + list_of_logprobs: list[dict[str, torch.Tensor]], + *, + original_data: BatchedDataDict[Any], + transformed_data: BatchedDataDict[Any], + metadata: dict[str, Any], + ) -> torch.Tensor: + """Scatter per-microbatch outputs back into one ``[B, S]`` tensor.""" + + # ---- iterator wrapping ------------------------------------------------- + + def _set_asymmetric_ar_metadata(self, microbatch: ProcessedMicrobatch) -> None: + """Configure attention and RoPE positions for one replay microbatch.""" + data_dict = microbatch.data_dict + if "diffu_grpo_noisy_lengths" not in data_dict: + return + + noisy_lengths = data_dict["diffu_grpo_noisy_lengths"] + noisy_valid_lengths = data_dict["diffu_grpo_noisy_valid_lengths"] + clean_padded_lengths = data_dict["diffu_grpo_clean_padded_lengths"] + noisy_response_offsets = data_dict["diffu_grpo_noisy_response_offsets"] + if noisy_lengths.numel() == 0: + return + if not torch.all(noisy_lengths == noisy_lengths[0]): + raise ValueError("Diffusion noisy length must be constant per microbatch") + if not torch.all(clean_padded_lengths == clean_padded_lengths[0]): + raise ValueError("Diffusion clean length must be constant per microbatch") + if not torch.all(noisy_response_offsets == noisy_response_offsets[0]): + raise ValueError( + "Diffusion noisy response offset must be constant per microbatch" + ) + + noisy_length = int(noisy_lengths[0].item()) + clean_length = int(clean_padded_lengths[0].item()) + noisy_response_offset = int(noisy_response_offsets[0].item()) + prompt_lengths = data_dict["diffu_grpo_completion_starts"] + response_lengths = data_dict["diffu_grpo_response_lengths"] + clean_lengths = data_dict["diffu_grpo_clean_lengths"] + + processed_inputs = microbatch.processed_inputs + processed_inputs.position_ids = build_asymmetric_position_ids( + noisy_length=noisy_length, + clean_length=clean_length, + noisy_response_offset=noisy_response_offset, + prompt_lengths=prompt_lengths, + noisy_valid_lengths=noisy_valid_lengths, + ) + processed_inputs.model_kwargs["block_size"] = self._diffusion_block_size() + + modules = self._diffusion_attention_modules() + fallback_modules = [ + module + for module in modules + if not hasattr(module, "set_asymmetric_ar_metadata") + ] + unknown_modules = [ + module + for module in fallback_modules + if module.__class__.__name__ != "MinistralFlexAttention" + ] + if unknown_modules: + names = sorted({module.__class__.__name__ for module in unknown_modules}) + raise RuntimeError( + "Diffusion completion-only replay requires attention modules with " + "set_asymmetric_ar_metadata(), or the checked HF NLD compatibility " + f"adapter; unsupported modules: {names}" + ) + + block_mask = None + if fallback_modules: + block_mask = build_asymmetric_semi_ar_block_mask( + block_size=self._diffusion_block_size(), + noisy_length=noisy_length, + clean_length=clean_length, + noisy_response_offset=noisy_response_offset, + prompt_lengths=prompt_lengths, + noisy_valid_lengths=noisy_valid_lengths, + clean_lengths=clean_lengths, + ) + + for module in modules: + if hasattr(module, "set_asymmetric_ar_metadata"): + module.set_asymmetric_ar_metadata( + noisy_length=noisy_length, + clean_length=clean_length, + noisy_response_offset=noisy_response_offset, + prompt_lengths=prompt_lengths, + response_lengths=response_lengths, + noisy_valid_lengths=noisy_valid_lengths, + clean_lengths=clean_lengths, + ) + elif not install_hf_nld_asymmetric_mask(module, block_mask=block_mask): + raise RuntimeError( + f"Unable to install asymmetric attention on {type(module).__name__}" + ) + + def _clear_asymmetric_ar_metadata(self) -> None: + for module in self._diffusion_attention_modules(): + if hasattr(module, "clear_asymmetric_ar_metadata"): + module.clear_asymmetric_ar_metadata() + else: + clear_hf_nld_asymmetric_mask(module) + + def _wrap_diffusion_microbatch_iterator( + self, iterator: Iterator[ProcessedMicrobatch] + ) -> Iterator[ProcessedMicrobatch]: + try: + for microbatch in iterator: + self._set_asymmetric_ar_metadata(microbatch) + yield microbatch + finally: + self._clear_asymmetric_ar_metadata() + + def _wrap_training_microbatch_iterator( + self, iterator: Iterator[ProcessedMicrobatch], metadata: dict[str, Any] + ) -> Iterator[ProcessedMicrobatch]: + del metadata + return self._wrap_diffusion_microbatch_iterator(iterator) + + def _wrap_logprob_microbatch_iterator( + self, iterator: Iterator[ProcessedMicrobatch], metadata: dict[str, Any] + ) -> Iterator[ProcessedMicrobatch]: + del metadata + return self._wrap_diffusion_microbatch_iterator(iterator) + + # ---- helpers ----------------------------------------------------------- + + def _diffusion_block_size(self) -> int: + """Block size, taken from the model config so it cannot drift from the model.""" + overrides = self.cfg.get("hf_config_overrides") or {} + if "block_size" in overrides: + return int(overrides["block_size"]) + block_size = getattr(getattr(self, "model_config", None), "block_size", None) + if block_size is None: + raise ValueError( + "block_size is not resolvable from hf_config_overrides or the model " + "config. It must agree with the decode config or generation-KL will " + "silently blow up (the fork records 1.41 from a 32-vs-16 mismatch)." + ) + return int(block_size) + + def _dp_group(self) -> Any: + """DP process group -- the Automodel equivalent of parallel_state's.""" + return self.dp_mesh.get_group() diff --git a/tests/unit/models/automodel/test_automodel_setup.py b/tests/unit/models/automodel/test_automodel_setup.py index da18cc4158b..f7bb12b8fab 100644 --- a/tests/unit/models/automodel/test_automodel_setup.py +++ b/tests/unit/models/automodel/test_automodel_setup.py @@ -773,6 +773,7 @@ def test_setup_distributed_passes_correct_params( assert fsdp2_call_kwargs["sequence_parallel"] is False assert fsdp2_call_kwargs["activation_checkpointing"] is False assert fsdp2_call_kwargs["backend"] == "nccl" + assert fsdp2_call_kwargs["mp_policy"].cast_forward_inputs is True # Verify create_device_mesh was called with correct size params mesh_call_kwargs = mock_create_mesh.call_args[1] diff --git a/tests/unit/models/automodel/test_automodel_train.py b/tests/unit/models/automodel/test_automodel_train.py index be246d41c51..5ef73f5cff0 100644 --- a/tests/unit/models/automodel/test_automodel_train.py +++ b/tests/unit/models/automodel/test_automodel_train.py @@ -203,6 +203,23 @@ def test_forward_disallow_flash_attn_args( call_kwargs = mock_model.call_args[1] assert "flash_attn_kwargs" not in call_kwargs + def test_forward_passes_model_specific_kwargs( + self, mock_model, processed_inputs_no_flash + ): + processed_inputs_no_flash.model_kwargs["block_size"] = 16 + + model_forward(mock_model, processed_inputs_no_flash) + + assert mock_model.call_args.kwargs["block_size"] == 16 + + def test_forward_rejects_core_argument_override( + self, mock_model, processed_inputs_no_flash + ): + processed_inputs_no_flash.model_kwargs["input_ids"] = torch.zeros(1, 1) + + with pytest.raises(ValueError, match="cannot override core model arguments"): + model_forward(mock_model, processed_inputs_no_flash) + # ===================== # Test extract_logits diff --git a/tests/unit/models/automodel/test_diffusion_attention.py b/tests/unit/models/automodel/test_diffusion_attention.py new file mode 100644 index 00000000000..54e6b2c0ff5 --- /dev/null +++ b/tests/unit/models/automodel/test_diffusion_attention.py @@ -0,0 +1,105 @@ +# 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. + +import torch + +from nemo_rl.models.automodel.diffusion_attention import ( + asymmetric_semi_ar_mask_mod, + build_asymmetric_position_ids, + clear_hf_nld_asymmetric_mask, + install_hf_nld_asymmetric_mask, +) + + +def _dense_mask(mask_mod, batch_size: int, sequence_length: int) -> torch.Tensor: + dense = torch.zeros(batch_size, sequence_length, sequence_length, dtype=torch.bool) + for batch_idx in range(batch_size): + for query_idx in range(sequence_length): + for kv_idx in range(sequence_length): + dense[batch_idx, query_idx, kv_idx] = mask_mod( + torch.tensor(batch_idx), + torch.tensor(0), + torch.tensor(query_idx), + torch.tensor(kv_idx), + ) + return dense + + +def test_asymmetric_mask_uses_prompt_and_per_sample_valid_lengths(): + prompt_lengths = torch.tensor([2, 3]) + noisy_valid_lengths = torch.tensor([4, 2]) + clean_lengths = torch.tensor([6, 5]) + mask_mod = asymmetric_semi_ar_mask_mod( + block_size=2, + noisy_length=4, + noisy_response_offset=0, + prompt_lengths=prompt_lengths, + noisy_valid_lengths=noisy_valid_lengths, + clean_lengths=clean_lengths, + ) + mask = _dense_mask(mask_mod, batch_size=2, sequence_length=10) + + # Sample 0, first noisy block: its own noisy block + the two-token prompt. + assert torch.equal( + torch.nonzero(mask[0, 0], as_tuple=False).flatten(), + torch.tensor([0, 1, 4, 5]), + ) + # Second noisy block also sees the previous clean response block. + assert torch.equal( + torch.nonzero(mask[0, 2], as_tuple=False).flatten(), + torch.tensor([2, 3, 4, 5, 6, 7]), + ) + # Sample 1 has only two valid noisy positions. Invalid query rows retain a + # self edge so flex attention never receives a fully masked row. + assert torch.equal( + torch.nonzero(mask[1, 2], as_tuple=False).flatten(), torch.tensor([2]) + ) + # Clean padding is excluded per sample (sample 1 clean length is five). + assert not mask[1, :9, 9].any() + + +def test_asymmetric_position_ids_align_noisy_response_with_clean_sequence(): + position_ids = build_asymmetric_position_ids( + noisy_length=4, + clean_length=6, + noisy_response_offset=0, + prompt_lengths=torch.tensor([2, 3]), + noisy_valid_lengths=torch.tensor([4, 2]), + ) + + assert torch.equal(position_ids[0], torch.tensor([2, 3, 4, 5, 0, 1, 2, 3, 4, 5])) + assert torch.equal(position_ids[1], torch.tensor([3, 4, 0, 0, 0, 1, 2, 3, 4, 5])) + + +def test_hf_nld_adapter_is_narrow_and_clears_cache(): + class MinistralFlexAttention: + def __init__(self): + self.sbd_block_diff_mask = None + self.block_size_orig = 16 + + def set_attention_mode(self, mode, block_size=None): + self.mode = mode + self.block_size = block_size + + module = MinistralFlexAttention() + sentinel = object() + assert install_hf_nld_asymmetric_mask(module, block_mask=sentinel) + assert module.sbd_block_diff_mask is sentinel + assert clear_hf_nld_asymmetric_mask(module) + assert module.sbd_block_diff_mask is None + + class UnknownAttention: + pass + + assert not install_hf_nld_asymmetric_mask(UnknownAttention(), block_mask=sentinel) diff --git a/tests/unit/models/automodel/test_diffusion_train.py b/tests/unit/models/automodel/test_diffusion_train.py new file mode 100644 index 00000000000..0844ff3ec08 --- /dev/null +++ b/tests/unit/models/automodel/test_diffusion_train.py @@ -0,0 +1,156 @@ +# 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 unittest.mock import MagicMock + +import pytest +import torch + +pytest.importorskip("nemo_automodel") + +from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.models.automodel.data import ProcessedInputs, ProcessedMicrobatch +from nemo_rl.models.automodel.diffusion_train import ( + DiffusionLogprobsPostProcessor, + DiffusionLossPostProcessor, + same_position_logprobs, +) +from nemo_rl.models.automodel.train import forward_with_post_processing_fn + + +def _processor_kwargs(sampling_params: TrainingSamplingParams) -> dict: + return { + "cfg": {}, + "device_mesh": MagicMock(), + "cp_mesh": MagicMock(), + "tp_mesh": MagicMock(), + "cp_size": 1, + "enable_seq_packing": False, + "sampling_params": sampling_params, + } + + +def test_same_position_logprobs_scores_target_at_same_position(): + logits = torch.tensor([[[8.0, 1.0, 0.0], [0.0, 7.0, 1.0], [0.0, 1.0, 6.0]]]) + target_ids = torch.tensor([[0, 1, 2]]) + + actual = same_position_logprobs(logits, target_ids) + expected = ( + torch.log_softmax(logits.float(), dim=-1) + .gather(-1, target_ids.unsqueeze(-1)) + .squeeze(-1) + ) + + assert torch.allclose(actual, expected) + assert torch.all(actual > -0.01) + + +def test_logprob_processor_applies_temperature_exactly_once(): + raw_logits = torch.tensor([[[4.0, 1.0, -2.0], [1.0, 5.0, 0.0]]]) + target_ids = torch.tensor([[0, 1]]) + sampling_params = TrainingSamplingParams(temperature=2.0) + post_processor = DiffusionLogprobsPostProcessor( + **_processor_kwargs(sampling_params) + ) + processed_inputs = ProcessedInputs( + input_ids=torch.tensor([[2, 2]]), + seq_len=2, + attention_mask=torch.ones(1, 2), + position_ids=torch.arange(2).unsqueeze(0), + ) + processed_microbatch = ProcessedMicrobatch( + data_dict=BatchedDataDict( + { + "input_ids": processed_inputs.input_ids, + "diffu_grpo_target_ids": target_ids, + } + ), + processed_inputs=processed_inputs, + original_batch_size=1, + original_seq_len=2, + ) + model = MagicMock(return_value=raw_logits.clone()) + + actual, _metrics, _ = forward_with_post_processing_fn( + model=model, + post_processing_fn=post_processor, + processed_mb=processed_microbatch, + sampling_params=sampling_params, + ) + expected = ( + torch.log_softmax(raw_logits.float() / 2.0, dim=-1) + .gather(-1, target_ids.unsqueeze(-1)) + .squeeze(-1) + ) + + assert torch.allclose(actual, expected) + + +def test_mask_token_exclusion_and_top_k_are_applied_before_softmax(): + logits = torch.tensor([[[9.0, 8.0, 1.0]]]) + target_ids = torch.tensor([[1]]) + sampling_params = TrainingSamplingParams(top_k=1, top_p=1.0, temperature=1.0) + + actual = same_position_logprobs( + logits, + target_ids, + exclude_token_id=0, + sampling_params=sampling_params, + ) + + assert torch.equal(actual, torch.zeros_like(actual)) + + +def test_loss_processor_uses_noisy_segment_and_valid_token_override(): + loss_fn = MagicMock() + loss_fn.compute_from_aligned_tensors.return_value = ( + torch.tensor(0.25), + {"num_valid_samples": 1}, + ) + valid_toks_override = torch.tensor(3.0) + post_processor = DiffusionLossPostProcessor( + loss_fn=loss_fn, + valid_toks_override=valid_toks_override, + dp_size=1, + **_processor_kwargs(TrainingSamplingParams()), + ) + logits = torch.randn(1, 5, 7) + data = BatchedDataDict( + { + "diffu_grpo_target_ids": torch.tensor([[1, 2, 3, 4, 5]]), + "diffu_grpo_noisy_lengths": torch.tensor([2]), + "diffu_grpo_loss_mask": torch.tensor([[1.0, 1.0, 0.0, 0.0, 0.0]]), + "sample_mask": torch.ones(1), + "advantages": torch.ones(1, 5), + "prev_logprobs": torch.zeros(1, 5), + "generation_logprobs": torch.zeros(1, 5), + } + ) + processed_inputs = ProcessedInputs(input_ids=torch.ones(1, 5).long(), seq_len=5) + + loss, metrics = post_processor( + logits=logits, + data_dict=data, + processed_inputs=processed_inputs, + global_valid_seqs=torch.tensor(1.0), + global_valid_toks=torch.tensor(99.0), + ) + + assert loss.item() == 0.25 + assert metrics["num_valid_samples"] == 1 + call = loss_fn.compute_from_aligned_tensors.call_args.kwargs + assert call["curr_logprobs"].shape == (1, 2) + assert call["token_mask"].shape == (1, 2) + assert call["global_valid_toks"] is valid_toks_override