bridge.data.packing.offline#

Offline materialization of packed GPT SFT artifacts.

Module Contents#

Functions#

_get_shared_dataset_item

_init_shared_dataset_worker

_materialize_dataset_items

_pre_pad_data_point

Pad a single data point so its runtime segment length is divisible by the requested multiple.

tokenize_dataset

Tokenizes a dataset from the provided path using the specified tokenizer and prepares it for further processing.

prepare_gpt_sft_packed_data

Prepares a packed sequence dataset from a given input file and saves it to an output file.

Data#

logger

_shared_dataset

API#

bridge.data.packing.offline.logger#

‘getLogger(…)’

bridge.data.packing.offline._shared_dataset#

None

bridge.data.packing.offline._get_shared_dataset_item(i)#
bridge.data.packing.offline._init_shared_dataset_worker(dataset)#
bridge.data.packing.offline._materialize_dataset_items(dataset, num_workers)#
bridge.data.packing.offline._pre_pad_data_point(
data: dict,
max_seq_length: int,
max_stored_length_to_pad: int,
pad_id: int,
) → None#

Pad a single data point so its runtime segment length is divisible by the requested multiple.

Pads input_ids/context_ids with pad_id and loss_mask with 0 (no loss on pad positions). The chat preprocessing path (_chat_preprocess) returns torch tensors rather than plain lists, so values are normalized to lists before concatenating; this avoids a TypeError from tensor + list and keeps loss_mask the same length as input_ids so that grouped samples do not produce a ragged array in fill_packing_strategy.

Parameters:

  • data – A single tokenized example. Mutated in place.

  • max_seq_length – Hard upper bound for the runtime sequence length after next-token shifting.

  • max_stored_length_to_pad – Stored target length to pad/truncate to. This is the divisible runtime target plus one token because packed SFT labels are derived by shifting input_ids.

  • pad_id – Token id used to pad input_ids/context_ids.

bridge.data.packing.offline.tokenize_dataset(
path: pathlib.Path,
tokenizer: megatron.bridge.training.tokenizers.tokenizer.MegatronTokenizer,
max_seq_length: int,
seed: int,
dataset_kwargs: dict | None = None,
pad_seq_to_mult: int | None = 1,
num_tokenizer_workers: int = -1,
*,
dataset_builder: collections.abc.Callable[..., Any],
)#

Tokenizes a dataset from the provided path using the specified tokenizer and prepares it for further processing.

Parameters:

  • path (Path) – Path to the dataset file.

  • tokenizer (MegatronTokenizer) – The tokenizer to use for tokenization.

  • max_seq_length (int) – Maximum sequence length for the tokens.

  • seed (int) – Random seed for shuffling the dataset.

  • dataset_kwargs (dict | None) – Additional GPT SFT dataset construction options. Can include ‘chat’, ‘use_hf_tokenizer_chat_template’, ‘tool_schemas’, etc.

  • pad_seq_to_mult (int | None) – Optional multiple to pad each sequence to during packing preparation (e.g., set to 2 * context_parallel_size for THD CP).

  • num_tokenizer_workers – Number of worker processes used to materialize tokenized samples. Values less than or equal to 1 run serially.

  • dataset_builder – Builder-owned callable that constructs one unpacked GPT SFT split.

Returns:

A NumPy array containing the tokenized data.

Return type:

np.ndarray

bridge.data.packing.offline.prepare_gpt_sft_packed_data(
input_path: pathlib.Path,
output_path: pathlib.Path,
output_metadata_path: pathlib.Path,
packed_sequence_size: int,
tokenizer: megatron.bridge.training.tokenizers.tokenizer.MegatronTokenizer,
max_seq_length: int,
seed: int | None = 0,
packing_algorithm: str = 'first_fit_shuffle',
dataset_kwargs: dict | None = None,
pad_seq_to_mult: int | None = 1,
num_tokenizer_workers: int = -1,
*,
dataset_builder: collections.abc.Callable[..., Any],
)#

Prepares a packed sequence dataset from a given input file and saves it to an output file.

Parameters:

  • input_path (Path) – Path to the input dataset file.

  • output_path (Path) – Path to save the packed sequence data.

  • output_metadata_path (Path) – Path to save packing metadata.

  • packed_sequence_size (int) – The maximum size for each packed sequence.

  • tokenizer (MegatronTokenizer) – The tokenizer to use for tokenization.

  • max_seq_length (int) – Maximum sequence length for the tokens.

  • seed (int | None) – Random seed for shuffling (optional).

  • packing_algorithm (str) – The algorithm used for packing sequences currently supports “first_fit_shuffle” and “first_fit_decreasing”.

  • dataset_kwargs (dict | None) – Additional GPT SFT dataset construction options. Enables packing with chat templates, tool schemas, etc.

  • pad_seq_to_mult (int | None) – Optional multiple to pad each sequence to during packing preparation (e.g., set to 2 * context_parallel_size for THD CP).

  • num_tokenizer_workers – Number of worker processes used to materialize tokenized samples. Values less than or equal to 1 run serially.

  • dataset_builder – Builder-owned callable that constructs one unpacked GPT SFT split.

Returns:

Saves the packed sequence data to the specified output path.

Return type:

None