Amazon S3 Connector for PyTorch

The dcp package provides benchmarks to evaluate the performance of PyTorch's Distributed Checkpointing (DCP) feature when compared to the s3torchconnector library. These benchmarks use Fully Sharded Data Parallel (FSDP) to measure how efficiently model parameters, sharded across multiple GPUs/processes, are saved and loaded.

Benchmark Metrics

Save Benchmarks measure:

• Checkpoint saving throughput (MiB/s)
• Checkpoint "corrected" save durations (seconds), excluding model load duration on the device.

Load Benchmarks measure:

• Checkpoint loading throughput (MiB/s)
• Checkpoint "corrected" load durations (seconds), excluding process setup and model loading to device.

# Run save benchmarks (default behavior)
./utils/run_dcp_fsdp_benchmarks.sh

# Explicitly run save benchmarks
./utils/run_dcp_fsdp_benchmarks.sh --save

# Run load benchmarks
./utils/run_dcp_fsdp_benchmarks.sh --load

The connector supports three reader types via S3ReaderConstructor. Choose the one that best fits your access pattern:

1. Sequential Reader (S3ReaderConstructor.sequential()):

   • Use case: Processing entire objects or when repeated access is required.
   • Behavior: Downloads and buffers the entire S3 object in memory. Prioritizes performance.

2. Range-based Reader (S3ReaderConstructor.range_based()):

   • Use case: Large objects (100MB+) requiring sparse partial reads or in memory-constrained environments.
   • Behavior: Uses byte-range requests. Features adaptive buffering: small reads (< buffer_size) use an internal buffer, while large reads (≥ buffer_size) bypass it for direct transfer.

3. DCP-Optimized Reader (S3ReaderConstructor.dcp_optimized()):

   • Use case: PyTorch Distributed Checkpoint (DCP) loading for highest performance and memory efficiency.
   • Behavior: Default for S3StorageReader. Provides per-item buffers, zero-copy management, and range coalescing for selective data fetching. Requires sequential access patterns.

Note: S3Reader instances are not thread-safe. For DataLoader multiprocessing, each worker creates its own instance automatically.

To prevent S3 throttling (503 Slow Down errors) during high-throughput distributed training, S3StorageWriter supports different prefix strategies to distribute keys across S3 partitions.

Available Strategies:

1. RoundRobinPrefixStrategy: Distributes checkpoints across a list of user-provided prefixes in a round-robin fashion.
2. BinaryPrefixStrategy: Generates binary (base-2) prefixes for optimal partitioning.
3. HexPrefixStrategy: Uses hexadecimal (base-16) prefixes for a balance of efficiency and readability.

Custom Strategies

You can implement a custom strategy by extending S3PrefixStrategyBase and overriding the generate_prefix(self, rank: int) -> str method.

from s3torchconnector.dcp import RoundRobinPrefixStrategy, S3StorageWriter

# Example: RoundRobin
strategy = RoundRobinPrefixStrategy(
    user_prefixes=["shard1", "shard2", "shard3"],
    epoch_num=5
)

writer = S3StorageWriter(
    region=REGION,
    path="s3://my-bucket/checkpoints/",
    prefix_strategy=strategy
)

# Example: Custom Strategy
from s3torchconnector.dcp import S3PrefixStrategyBase

class CustomPrefixStrategy(S3PrefixStrategyBase):
    def __init__(self, custom_param):
        super().__init__()
        self.custom_param = custom_param

    def generate_prefix(self, rank: int) -> str:
        return f"custom_{self.custom_param}/{rank}/"

Benchmarks are executed by editing Hydra configuration files in the conf/ directory (ensuring all ??? placeholders are replaced) and then running the corresponding shell script.

Important: Ensure your S3 bucket is in the same region as your EC2 instance to avoid network latency affecting measurements. You can override Hydra parameters via the command line using the + prefix.

Available Scenarios

# Example: Running dataset benchmarks with a parameter override
./utils/run_dataset_benchmarks.sh +disambiguator=some_key

To run benchmarks, it is recommended to use an EC2 instance with GPUs (e.g., G5 instance type) using an AWS Deep Learning AMI (DLAMI). For best results, use a dedicated instance without other resource-intensive processes. To avoid token expiration during long-running benchmarks, attach an IAM role to the EC2 instance with:

• Full access to S3
• (Optional) Full access to DynamoDB (required if writing run results to DynamoDB)

Follow these steps to create a clean Conda environment and install the package:

1. Create a new Conda environment with Python 3.12.
2. Initialize Conda and activate the environment.
3. Install the s3torchbenchmarking module in editable mode using pip install -e ..

Note: Some scenarios may require the Mountpoint for Amazon S3 file client.

# Create environment
conda create -n pytorch-benchmarks python=3.12
conda init

# Activate (you may need to log out/in after conda init)
source activate pytorch-benchmarks

# Install the module
pip install -e .

While Python logging handles the Python-side implementation, the Rust components use tracing_subscriber. You can control Rust-side logs using the S3_TORCH_CONNECTOR_DEBUG_LOGS environment variable.

Note on AWS CRT logs: AWS CRT logs are very noisy. To avoid noise, append awscrt=off to your log configuration.

To save logs to a file instead of STDOUT, set the S3_TORCH_CONNECTOR_LOGS_DIR_PATH environment variable to a writable directory. Logs will be written to ${S3_TORCH_CONNECTOR_LOGS_DIR_PATH}/s3torchconnectorclient.log.yyyy-MM-dd-HH and rolled hourly.

# Configure INFO level logs to be written to STDOUT
export S3_TORCH_CONNECTOR_DEBUG_LOGS=info

# Enable TRACE level logs (most verbose) written to a directory
export S3_TORCH_CONNECTOR_DEBUG_LOGS=trace
export S3_TORCH_CONNECTOR_LOGS_DIR_PATH="/tmp/s3torchconnector-logs"

# Enable TRACE level logs with AWS CRT logs filtered out
export S3_TORCH_CONNECTOR_DEBUG_LOGS=trace,awscrt=off
export S3_TORCH_CONNECTOR_LOGS_DIR_PATH="/tmp/s3torchconnector-logs"

# Set up different levels for inner components
export S3_TORCH_CONNECTOR_DEBUG_LOGS=trace,mountpoint_s3_client=debug,awscrt=error