Implementation:Microsoft DeepSpeedExamples Tensor Parallel Bench Length Padding

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	Tensor Parallelism, Benchmarking
Last Updated	2026-02-07 12:00 GMT

Overview

An alternative tensor parallelism training script that pads all sequences to max_length for benchmarking uniform batch sizes, with pad token masking in labels.

Description

This script is a variant of the standard tensor parallel training script, specifically modified for benchmarking purposes by padding all tokenized sequences to max_length instead of using dynamic longest padding. The key difference is in the _tokenize_fn function, which uses padding="max_length" to ensure every sequence in a batch has the same length, eliminating variable-length overhead and providing consistent timing measurements.

The preprocess function includes an additional post-processing step that replaces pad token IDs in the labels with IGNORE_INDEX (-100), ensuring that padding tokens do not contribute to the loss computation. This is necessary because max_length padding introduces trailing pad tokens that must be excluded from the cross-entropy loss calculation.

All other components are identical to the standard training script: ModelArguments, DataArguments, and TrainingArguments dataclasses, smart_tokenizer_and_embedding_resize for special token handling, SupervisedDataset with pickle caching (using a length-specific cache filename dataset_dict{max_length}.pkl), DataCollatorForSupervisedDataset for batch collation, and the train function with MemoryCallback for GPU/CPU memory monitoring. The dataset cache filename incorporates the model_max_length to avoid conflicts between runs with different sequence lengths.

Usage

Use this script specifically for benchmarking tensor parallel training throughput, where consistent sequence lengths across batches are required to produce reliable performance measurements. It should not be used for production fine-tuning since max_length padding wastes compute on padding tokens.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/tensor_parallel/train_bench_length.py
• Lines: 1-274

Signature

@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")

@dataclass
class DataArguments:
    data_path: str = field(default=None)

@dataclass
class TrainingArguments(transformers.TrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    optim: str = field(default="adamw_torch")
    model_max_length: int = field(default=512)

def smart_tokenizer_and_embedding_resize(
    special_tokens_dict: Dict,
    tokenizer: transformers.PreTrainedTokenizer,
    model: transformers.PreTrainedModel,
):
    ...

def _tokenize_fn(strings: Sequence[str],
                 tokenizer: transformers.PreTrainedTokenizer) -> Dict:
    ...  # Uses padding="max_length"

def preprocess(
    sources: Sequence[str],
    targets: Sequence[str],
    tokenizer: transformers.PreTrainedTokenizer,
) -> Dict:
    ...  # Includes pad token masking in labels

class SupervisedDataset(Dataset):
    def __init__(self, data_path: str, tokenizer: transformers.PreTrainedTokenizer):
        ...

@dataclass
class DataCollatorForSupervisedDataset(object):
    tokenizer: transformers.PreTrainedTokenizer
    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        ...

def train():
    ...

Import

from train_bench_length import SupervisedDataset, DataCollatorForSupervisedDataset, train

I/O Contract

Inputs

Name	Type	Required	Description
--model_name_or_path	str	No	HuggingFace model identifier or path (default: 'facebook/opt-125m')
--data_path	str	Yes	Path to the JSON training data file in Alpaca format
--cache_dir	str	No	Cache directory for model and tokenizer downloads
--optim	str	No	Optimizer name (default: 'adamw_torch')
--model_max_length	int	No	Maximum sequence length; all sequences are padded to this length (default: 512)
--output_dir	str	Yes	Directory to save the trained model

Outputs

Name	Type	Description
saved model	directory	HuggingFace model checkpoint saved to output_dir
dataset_dict{max_length}.pkl	file	Cached tokenized dataset keyed by max_length for faster reloading
stdout	text	Memory usage statistics after each training step

Usage Examples

# Benchmark with fixed sequence length of 1024
# deepspeed train_bench_length.py --model_name_or_path meta-llama/Llama-2-7b-hf \
#     --data_path ./alpaca_data.json \
#     --output_dir ./bench_output \
#     --model_max_length 1024 \
#     --deepspeed ds_config.json

# Compare throughput at different sequence lengths
# deepspeed train_bench_length.py --model_max_length 512 ...
# deepspeed train_bench_length.py --model_max_length 1024 ...
# deepspeed train_bench_length.py --model_max_length 2048 ...