From eab11e0117d701646f10360ef5d9ecee6fa35b9b Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Sun, 19 Jul 2026 21:14:00 -0500 Subject: [PATCH 1/2] Skip attention mask materialization in RL training Signed-off-by: Teodor-Dumitru Ene --- megatron/rl/rl_utils.py | 7 +++++-- megatron/training/utils/common_utils.py | 23 +++++++++++++++-------- train_rl.py | 4 ++++ 3 files changed, 24 insertions(+), 10 deletions(-) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 0b834def67b..01ba0f4fcf3 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -1704,16 +1704,19 @@ def prepare_data_for_update( data_loader = DataLoader(dataset, batch_size=1) logprobs_batch_size = 1 else: - # Always compute standard masks for the original data (we'll need them later) + # Compute the loss mask and position ids for the original data (we'll need them later). + # No dense attention mask: the forward pass masks via PackedSeqParams (see + # get_logprobs), even when sequence packing is disabled. with nvtx_range("rl/get-ltor-masks", time=True): _, original_loss_mask, original_position_ids = get_ltor_masks_and_position_ids( trajs, tokenizer.eod, tokenizer.pad, args.reset_position_ids, - args.reset_attention_mask, + reset_attention_mask=False, eod_mask_loss=False, pad_mask_loss=True, + create_attention_mask=False, ) original_loss_mask[~generation_masks] = 0.0 compute_trajs = trajs diff --git a/megatron/training/utils/common_utils.py b/megatron/training/utils/common_utils.py index 30617ef9b4c..f65d5b89972 100644 --- a/megatron/training/utils/common_utils.py +++ b/megatron/training/utils/common_utils.py @@ -377,20 +377,26 @@ def get_ltor_masks_and_position_ids(data, reset_position_ids, reset_attention_mask, eod_mask_loss, - pad_mask_loss): + pad_mask_loss, + create_attention_mask=True): """Build masks and position id for left to right model.""" + assert create_attention_mask or not reset_attention_mask, \ + "reset_attention_mask requires the attention mask to be created." # Extract batch size and sequence length. micro_batch_size, seq_length = data.size() # Attention mask (lower triangular). - if reset_attention_mask: - att_mask_batch = micro_batch_size + if create_attention_mask: + if reset_attention_mask: + att_mask_batch = micro_batch_size + else: + att_mask_batch = 1 + attention_mask = torch.tril( + torch.ones((att_mask_batch, seq_length, seq_length), device=data.device) + ).view(att_mask_batch, 1, seq_length, seq_length) else: - att_mask_batch = 1 - attention_mask = torch.tril( - torch.ones((att_mask_batch, seq_length, seq_length), device=data.device) - ).view(att_mask_batch, 1, seq_length, seq_length) + attention_mask = None # Loss mask. loss_mask = torch.ones(data.size(), dtype=torch.float, device=data.device) @@ -429,7 +435,8 @@ def get_ltor_masks_and_position_ids(data, prev_index = i + 1 # Convert attention mask to binary: - attention_mask = attention_mask < 0.5 + if create_attention_mask: + attention_mask = attention_mask < 0.5 return attention_mask, loss_mask, position_ids diff --git a/train_rl.py b/train_rl.py index acf54680f4a..17eec357b15 100644 --- a/train_rl.py +++ b/train_rl.py @@ -415,6 +415,10 @@ def _model_builder( extra_args_provider=add_inference_args, args_defaults={}, ) + assert not args.reset_attention_mask, ( + "--reset-attention-mask is not supported in RL training: " + "the forward pass masks via PackedSeqParams and never consumes a dense attention mask." + ) if is_hybrid_model(args): model_cfg = hybrid_config_from_args(args) else: From d10492096a98ea48a9ff74ee79ace496eccb7820 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 12 Aug 2026 09:48:18 -0500 Subject: [PATCH 2/2] Address reviewer comments Signed-off-by: Teodor-Dumitru Ene --- megatron/training/utils/common_utils.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/megatron/training/utils/common_utils.py b/megatron/training/utils/common_utils.py index f65d5b89972..1b20eb5a985 100644 --- a/megatron/training/utils/common_utils.py +++ b/megatron/training/utils/common_utils.py @@ -379,7 +379,25 @@ def get_ltor_masks_and_position_ids(data, eod_mask_loss, pad_mask_loss, create_attention_mask=True): - """Build masks and position id for left to right model.""" + """Build masks and position id for left to right model. + + Args: + data: Token ids, shape [micro_batch_size, seq_length]. + eod_token: End-of-document token id. + pad_token: Padding token id. + reset_position_ids: Restart position ids from 0 after each EOD token. + reset_attention_mask: Additionally mask attention across document boundaries, + turning the shared causal mask into a per-sample block-causal mask. + Requires create_attention_mask, since it modifies the materialized mask. + eod_mask_loss: Zero the loss mask at EOD tokens. + pad_mask_loss: Zero the loss mask at pad tokens. + create_attention_mask: Materialize the dense causal attention mask. + Can be disabled if the attention kernel generates the mask by itself + (e.g. from PackedSeqParams), in which case attention_mask is returned as None. + + Returns: + Tuple of (attention_mask or None, loss_mask, position_ids). + """ assert create_attention_mask or not reset_attention_mask, \ "reset_attention_mask requires the attention mask to be created."