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..1b20eb5a985 100644 --- a/megatron/training/utils/common_utils.py +++ b/megatron/training/utils/common_utils.py @@ -377,20 +377,44 @@ def get_ltor_masks_and_position_ids(data, reset_position_ids, reset_attention_mask, eod_mask_loss, - pad_mask_loss): - """Build masks and position id for left to right model.""" + pad_mask_loss, + create_attention_mask=True): + """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." # 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 +453,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: