`bridge.data.packing.gpt_sft`#

Runtime dataset for legacy NumPy GPT SFT packed artifacts.

Module Contents#

Classes#

GPTSFTPackedDataset

Functions#

_safe_load_packed_npy

Load a packed NumPy artifact with the restricted object unpickler.

Data#

logger

API#

bridge.data.packing.gpt_sft.logger#: ‘getLogger(…)’

bridge.data.packing.gpt_sft._safe_load_packed_npy(file_path: str | pathlib.Path) → numpy.ndarray#: Load a packed NumPy artifact with the restricted object unpickler.

class bridge.data.packing.gpt_sft.GPTSFTPackedDataset(

file_path: str,

tokenizer: megatron.bridge.training.tokenizers.tokenizer.MegatronTokenizer,

return_cu_seqlen: bool = True,

pad_cu_seqlens: bool = False,

pad_seq_to_mult: int = 1,

pack_metadata_file_path: str | None = None,

**kwargs,

)#

Bases: megatron.bridge.data.datasets.gpt_sft.GPTSFTDataset

Initialization

file_path: See file_path in the parent class. tokenizer: See tokenizer in the parent class. return_cu_seqlen: Whether to return cu_seqlen to pass to the model. Having cu_seqlen in the model input enables THD attention kernel, which is the correct format for training with packed sequence to prevent cross-sequence attention. This flag should be True unless you have a specific use case. pad_seq_to_mult: The multiple used for padding sequences during packing. When > 1, cu_seqlens_unpadded will be computed to support THD CP. When == 1 (no padding), cu_seqlens_unpadded is not computed.

__getitem__(idx)#

_load_dataset()#

_build_samples_mapping()#

_build_loss_mask(processed_example)#

_maybe_cast_to_list(x)#

collate_fn(batch)#

Collates a list of packed sequence samples into a batch for the model.

This method is specifically designed for GPTSFTPackedDataset. It takes a list of packed sequence items (as returned by __getitem__) and prepares a batch of tensors. This includes handling cu_seqlens which are crucial for the efficient processing of packed sequences with kernels like THD attention.

Parameters:: batch (List[dict]) – A list of packed sequence samples.
Returns:: A dictionary of batched tensors, including ‘tokens’, ‘labels’, ‘loss_mask’, ‘position_ids’, and potentially ‘cu_seqlens’, ‘cu_seqlens_argmin’, ‘max_seqlen’ if return_cu_seqlen is True.
Return type:: dict

bridge.data.packing.gpt_sft#