Problem
I finetuine a model and set cp > 1，it raise a error, the error val is attention_mask , it is [1], it's shape is invalid.
51: [rank51]: File "/opt/Megatron-Bridge/src/megatron/bridge/training/gpt_step.py", line 207, in get_batch
51: [rank51]: batch = _partition_packed_batch_for_cp(batch, cp_size)
51: [rank51]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
51: [rank51]: File "/opt/Megatron-Bridge/src/megatron/bridge/training/gpt_step.py", line 98, in _partition_packed_batch_for_cp
51: [rank51]: index = tex.thd_get_partitioned_indices(cu_seqlens, val.size(1), cp_size, cp_rank)
51: [rank51]: ^^^^^^^^^^^
51: [rank51]: IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)
Minimal repro
I find attention_mask is added at here:
if include_full_batch_fields:
required_device_keys.update(("tokens", "labels", "loss_mask", "attention_mask", "position_ids"))
elif not skip_getting_attention_mask_from_dataset:
required_device_keys.add("attention_mask")
if "cu_seqlens" in batch:
required_device_keys.add("cu_seqlens")
if "cu_seqlens_unpadded" in batch:
required_device_keys.add("cu_seqlens_unpadded")
required_host_keys.add("cu_seqlens_argmin")
required_host_keys.add("max_seqlen")
if "cu_seqlens_unpadded_argmin" in batch:
required_host_keys.add("cu_seqlens_unpadded_argmin")
Expected behavior
I think the correct code is like this:
if include_full_batch_fields:
required_device_keys.update(("tokens", "labels", "loss_mask", "attention_mask", "position_ids"))
elif not skip_getting_attention_mask_from_dataset:
required_device_keys.add("attention_mask")
if "cu_seqlens" in batch:
required_device_keys.add("cu_seqlens")
if "cu_seqlens_unpadded" in batch:
required_device_keys.add("cu_seqlens_unpadded")
required_host_keys.add("cu_seqlens_argmin")
required_host_keys.add("max_seqlen")
if "cu_seqlens_unpadded_argmin" in batch:
required_host_keys.add("cu_seqlens_unpadded_argmin")
# Packed sequence data does not require an attention mask <---- should remove it
if 'attention_mask' in required_device_keys:
required_device_keys.remove('attention_mask')
Affected area
area:training
Regression?
Yes
Environment
No response
Logs
Problem
I finetuine a model and set cp > 1，it raise a error, the error val is attention_mask , it is [1], it's shape is invalid.
Minimal repro
I find attention_mask is added at here:
Expected behavior
I think the correct code is like this:
Affected area
area:training
Regression?
Yes
Environment
No response
Logs