New in Composer 0.11: FSDP Support, Streaming v0.1 Release, Simplified Checkpointing and Distributed Experience

import torch.nn as nn
from composer import Trainer

class Block (nn.Module):
    ...

# Your custom model
class Model(nn.Module):
    def __init__(self, n_layers):
        super().__init__()
        self.blocks = nn.ModuleList([
            Block(...) for _ in range(n_layers)
        ]),
        self.head = nn.Linear(...)
    def forward(self, inputs):
        ...

    # FSDP Wrap Function
    def fsdp_wrap_fn(self, module):
        return isinstance(module, Block)

    # Activation Checkpointing Function
    def activation_checkpointing_fn(self, module):
        return isinstance(module, Block)

# ComposerModel wrapper, used by the Trainer
# to compute loss, metrics, etc.
class MyComposerModel(ComposerModel):

    def __init__(self, n_layers):
        super().__init__()
        self.model = Model(n_layers)
        ...

    def forward(self, batch):
        ...

    def eval_forward(self, batch, outputs=None):
        ...

    def loss(self, outputs, batch):
        ...
		    
# Pass your ComposerModel and fsdp_config into the Trainer
composer_model = MyComposerModel(n_layers=3)
fsdp_config = {
    'sharding_strategy': 'FULL_SHARD',
    'min_params': 1e8,
    'cpu_offload': False, # Not supported yet
    'mixed_precision': 'DEFAULT',
    'backward_prefetch': 'BACKWARD_POST',
    'activation_checkpointing': False,
    'activation_cpu_offload': False,
    'verbose': True
}

trainer = Trainer(
    model=composer_model,
    fsdp_config=fsdp_config,
    ...
)

trainer.fit()

pip install mosaicml-streaming

import torch
from streaming import Dataset

dataloader = torch.utils.data.DataLoader(dataset=Dataset(remote='s3://...'))

import torch
from streaming.base import Dataset

# Extending `streaming.Dataset` with custom get functionality
class CustomDataset(Dataset):
   def __init__(self, local, remote):
       super().__init__(local, remote)

   def __getitem__(self, idx: int) -> Any:
       obj = super().__getitem__(idx)
       return obj['x'], obj['y']

# Local caching directory
local = '/tmp/cache'

# Remote location to stream from
remote ='s3://mybucket/myfolder'

dataloader = torch.utils.data.DataLoader(dataset=CustomDataset(local=local, remote=remote))

• Specify with save_folder your full URI to your save directory destination (e.g. 's3://my-bucket/{run_name}/checkpoints')
• Optionally, set save_filename to the pattern you want for your checkpoint file names

from composer.trainer import Trainer

# Checkpoint saving to S3.
trainer = Trainer(
    model=model,
    save_folder="s3://my-bucket/{run_name}/checkpoints",
    run_name='my-run',
    save_interval="1ep",
    save_filename="ep{epoch}.pt",
    save_num_checkpoints_to_keep=0,  # delete all checkpoints locally
    ...
)

trainer.fit()

• Set load_path to the full URI to your desired checkpoint file (e.g. 's3://my-bucket/my-run/checkpoints/epoch13.pt')

from composer.trainer import Trainer

# Checkpoint loading from S3.
new_trainer = Trainer(
    model=model,
    train_dataloader=train_dataloader,
    max_duration="10ep",
    load_path="s3://my-bucket/my-run/checkpoints/ep13.pt",
    ...
 )

new_trainer.fit()

import datetime
from composer.trainer.devices import DeviceGPU
from composer.utils import dist

dist.initialize(DeviceGPU(), datetime.timedelta(seconds=30)) # Initialize distributed module

if dist.get_local_rank() == 0: # Download dataset on rank zero
    dataset = download_my_dataset()
dist.barrier() # All ranks wait until dataset is downloaded

# Create and train your model!