Day 4 - Optimizing Training Data Storage and Access: Pipe Mode vs File Mode, FSx for Lustre, and Distributed Training

📌 Key Takeaways

• Even with well-prepared data, if reads are slow the GPU sits idle. I/O bottlenecks are the most expensive waste in practice.
• The tradeoff among File (download everything, then start) / Pipe (streaming, starts immediately) / FastFile (streaming + POSIX access) is startup delay vs random-access flexibility.
• FSx for Lustre pays off when the workload is large + repeatedly read + low latency + distributed. It is overkill for one-off, small jobs.
• Distributed training is split into data parallelism (replicate the model, split the data) and model parallelism (split the model itself).
• ShardedByS3Key only works when the data is split across multiple S3 objects. It has no effect on a single giant file.

SageMaker Input Modes

There are several ways for a SageMaker training job to bring data from S3 into the container. "How do you stream the data into the training container" directly affects performance and cost.

File Mode (Default)

Downloads all of the S3 data to the container's EBS volume before training starts, and only then begins training.

• Pros: Simple and stable. You handle the data like local files. Good for reading repeatedly over multiple epochs.
• Cons: If the data is large, training does not start until the download finishes (startup delay). The EBS capacity must be at least the size of the dataset.

Pipe Mode

Instead of downloading in advance, it streams directly from S3 during training.

• Pros: There is no download wait, so training starts immediately. It is free of disk capacity limits and can handle datasets larger than the disk. Throughput is high on large data.
• Cons: Because it is a sequential stream, random access is hard, and it usually requires a specific format such as RecordIO/protobuf.

FastFile Mode

A relatively newer mode that is a compromise between File mode's ease of use (access like POSIX files) and Pipe mode's immediate start. It fetches files by streaming when needed but lets you treat them like local files. Because it allows random access without waiting for a full download, it is often recommended for large datasets.

# Specify the input mode in a SageMaker Estimator (conceptual example)
from sagemaker.estimator import Estimator
from sagemaker.inputs import TrainingInput
 
estimator = Estimator(
    image_uri=image,
    role=role,
    instance_count=1,
    instance_type="ml.p3.2xlarge",
    input_mode="Pipe",   # "File" | "Pipe" | "FastFile"
)
 
train_input = TrainingInput(
    s3_data="s3://ml-train/clicks/",
    input_mode="FastFile",   # the mode can be specified again per input
)
estimator.fit({"train": train_input})

File mode is "download everything, then start," and Pipe mode is "start while downloading." The key tradeoff is startup delay vs random-access flexibility. If the data is hundreds of GB to TB and you worry about startup delay and disk cost, Pipe is better; if the data is small and read often and flexibly, File is better.

⚠️ Trap: The simplification "Pipe mode is always faster" is dangerous. Pipe is a sequential stream, so random access and shuffling are difficult and it requires a specific format. For small data or data that needs random access, File/FastFile is actually the better fit. You must weigh data size, access pattern, and format together.

FSx for Lustre — An Ultra-High-Performance Training File System

When you want to cut the latency of direct S3 access even further for large-scale, repeated training, use Amazon FSx for Lustre. Lustre is a parallel file system for HPC (high-performance computing) that provides very high throughput and low latency.

Key features:

• S3 integration: When you link FSx for Lustre to an S3 bucket, the data is loaded and synchronized automatically. S3 acts as the durable store, and FSx acts as a high-speed cache and workspace.
• Strong for repeated training: When the same data is read repeatedly across multiple epochs and multiple jobs, it is read quickly from FSx instead of being fetched from S3 each time.
• Shared storage for distributed training: Multiple training nodes read the same file system at high speed at the same time.

The classic situation where FSx for Lustre shines is when (1) the dataset is large, (2) it is read repeatedly across multiple training jobs and epochs, and (3) low latency matters. For one-off or small training jobs, direct S3 access (File/FastFile) is enough, and FSx's extra cost and setup are not justified.

Data Sharding and Distributed Training

When the data or the model does not fit on one machine, you use distributed training. The key is how you divide the work across multiple nodes.

Data Parallelism vs Model Parallelism

[Data parallelism]
 Data    ┌── Shard 1 ──▶ GPU1 (model replica) ─┐
        ├── Shard 2 ──▶ GPU2 (model replica) ─┼─▶ Gradient sync ─▶ Weight update
        ├── Shard 3 ──▶ GPU3 (model replica) ─┤
        └── Shard 4 ──▶ GPU4 (model replica) ─┘

[Model parallelism]
 Batch ──▶ GPU1(layers 1-4) ──▶ GPU2(layers 5-8) ──▶ GPU3(layers 9-12) ──▶ Output

Sharding (ShardedByS3Key)

In data-parallel training, it is efficient for each node to read only its own share of the data. SageMaker controls this with S3DataDistributionType.

• FullyReplicated (default): Every instance receives the entire dataset.
• ShardedByS3Key: Splits the S3 objects (files) across instances so each gets only its own share. In data-parallel training this removes duplicate downloads and duplicate processing, improving efficiency.

# Specify sharding in distributed data-parallel training (conceptual example)
train_input = TrainingInput(
    s3_data="s3://ml-train/clicks/",
    distribution="ShardedByS3Key",   # each instance receives only part of the data
)
estimator = Estimator(
    image_uri=image, role=role,
    instance_count=4,                # distributed across 4 nodes
    instance_type="ml.p3.16xlarge",
)
estimator.fit({"train": train_input})

For ShardedByS3Key to be effective, the data must be split into multiple S3 objects (files). A single giant file cannot be split at the object level, so sharding has no effect at all. Therefore, for large data it is distributed-training friendly to pre-split it into multiple moderately sized files (for example, hundreds of Parquet/RecordIO shards).

⚠️ Trap: If you also split the validation set with ShardedByS3Key, each node sees different validation data and the metrics become inconsistent. The validation channel is usually correct with FullyReplicated.

When You Actually Turn On Distributed Training

Distribution is not free. Adding nodes adds communication cost, and speed does not increase linearly with the number of nodes.

• SageMaker distributed libraries: Separate libraries are provided for data parallelism and model parallelism, so you can turn on distribution without changing the framework code much.
• Batch size and learning rate: If you increase nodes by N times, the effective batch size also becomes N times. If you do not adjust the learning rate along with it, convergence gets worse.
• Communication bottleneck: Gradient synchronization happens at every step, so network bandwidth between nodes matters. If adding nodes actually makes training slower, suspect communication.
• Checkpoints: Long training runs must save intermediate checkpoints to S3. In particular, with managed Spot training you can resume from the checkpoint after an interruption.

A setup that runs long training on Spot instances without checkpoints starts over from the beginning when interrupted. Remember "cost savings" and "checkpoints" as a pair.

The Order for Diagnosing an I/O Bottleneck

When a scenario says "training is slow," suspect the supply side first, not the computation.

1. Is GPU utilization low? → If it is low, the bottleneck is the data supply, not the computation.
2. Does it take long to start? → Likely the wait for the full download in File mode. Consider Pipe/FastFile.
3. Is the same data read repeatedly? → Get the caching benefit with FSx for Lustre.
4. Are there too many small files? → Consolidate them into shards to reduce request overhead.
5. Is each node receiving everything? → Distribute with ShardedByS3Key.

🎯 Scenario: "We train data-parallel on several TB of data across 8 GPU nodes, but each node receives the entire dataset, so startup is too slow." → Split the data into multiple files and use ShardedByS3Key so each node receives only its own share. Duplicate downloads and duplicate processing disappear. If there are many repeated reads, attach FSx for Lustre as shared storage to cut latency further.

Instance Choice Is Also Part of Data Supply

The same data strategy gives different results depending on the instance type.

• GPU instances have a high hourly cost, so idle time is a direct loss. That is why I/O optimization is cost optimization.
• The standard approach is to separate them: run preprocessing as a separate job on CPU instances and use GPU instances only for training.

"Training is slow, so let's switch to a bigger GPU" is only right when GPU utilization is already high. If utilization is low and you only enlarge the GPU, the idle time just gets more expensive.

Tomorrow we tie all of Week 2 (transformation, pipelines, augmentation, storage optimization) together in a comprehensive review.

📖 Glossary

• Input mode: The way S3 data is brought into the training container. File, Pipe, FastFile.
• Startup delay: The time a training job spends preparing data before it actually begins computing.
• POSIX access: An access style where, as with a normal file system, you open by path and read at arbitrary positions.
• FSx for Lustre: A parallel file system for HPC. It uses S3 as the backend and behaves like a high-speed cache.
• Data parallelism: A method that replicates the model, splits the data for training, and then synchronizes the gradients.
• Model parallelism: A method that splits the model itself across multiple GPUs. Used for large models.
• Gradient synchronization: The process of combining the gradients each node computed so that all nodes have the same weights. It incurs communication cost.
• ShardedByS3Key: A distribution setting that splits S3 objects across instances so each receives only a different piece.
• FullyReplicated: The default distribution setting where every instance receives the entire data identically.
• Checkpoint: A saved snapshot of the model state during training. It lets you resume training after an interruption.