Skip to content

Quickstart

About

This page trains a small model on one GPU in a few seconds. The job has the same parts as every d9d job: a model, a model provider, a task, a data provider and a configuration. The code is in example/guide/quickstart.

Run the Example

Install d9d (see Installation) and PyYAML, which reads the config. Then get the example from the repository and run it:

1
2
3
4
pip install d9d pyyaml
git clone https://github.com/d9d-project/d9d.git
cd d9d/example/guide/quickstart
python train.py config.yaml

A job without parallelism runs as a single process, so it does not need torchrun. The job writes the loss and the gradient norm of every step to its log:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
...
[d9d] [local] ... - INFO - Starting job from scratch
[d9d] [local] ... - INFO - step 0 (stage=train): loss=13.2302
[d9d] [local] ... - INFO - step 0 (stage=train): l2_grad_norm_total=23.721
[d9d] [local] ... - INFO - step 1 (stage=train): loss=15.0928
[d9d] [local] ... - INFO - step 1 (stage=train): l2_grad_norm_total=23.8028
...
[d9d] [local] ... - INFO - step 255 (stage=train): loss=0.0134067
[d9d] [local] ... - INFO - step 255 (stage=train): l2_grad_norm_total=0.484559
Training: 100%|██████████| 256/256

The job writes its results to runs/:

  • runs/checkpoints/quickstart/: The two latest checkpoints, save-200 and save-256.
  • runs/export/: The trained weights as .safetensors files.

If you start the job again, it loads save-256 and logs Training is already complete, nothing to do. See Running Jobs for checkpoints and resuming.

The Code

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
import dataclasses
import sys
from collections.abc import Sequence
from pathlib import Path

import torch
import yaml
from d9d.core.dist_context import DENSE_DOMAIN, DeviceMeshParameters
from d9d.loop.auto import (
    AutoDataConfig,
    AutoDataProvider,
    AutoLRSchedulerConfig,
    AutoLRSchedulerProvider,
    AutoOptimizerConfig,
    AutoOptimizerProvider,
)
from d9d.loop.config import TrainerConfig
from d9d.loop.control import (
    BuildForwardInputsContext,
    BuildForwardInputsResult,
    ComputeLossContext,
    ComputeLossResult,
    InitializeModelStageContext,
    InitializeModelStageResult,
    ModelProvider,
    ParallelizeModelStageContext,
    PrepareExportModelStageContext,
    PrepareExportModelStageResult,
    TrainTask,
)
from d9d.loop.run import TrainingConfigurator
from d9d.model_state.mapper.adapters import identity_mapper_from_module
from d9d.module.parallelism.api import parallelize_hsdp
from pydantic import BaseModel, ConfigDict
from torch import nn
from torch.utils.data import Dataset


class ProjectConfig(BaseModel):
    # Rejects unknown keys, so a misspelled key in config.yaml fails at load.
    model_config = ConfigDict(extra="forbid")

    mesh: DeviceMeshParameters
    data: AutoDataConfig
    trainer: TrainerConfig
    optimizer: AutoOptimizerConfig
    lr_scheduler: AutoLRSchedulerConfig
    export_to: Path


# One sample of the dataset.
@dataclasses.dataclass
class RegressionSample:
    x: torch.Tensor
    y: torch.Tensor


# One microbatch, as collate() builds it.
@dataclasses.dataclass
class RegressionBatch:
    x: torch.Tensor
    y: torch.Tensor


# What the model receives.
@dataclasses.dataclass
class RegressionInput:
    x: torch.Tensor


# What compute_loss() needs but the model does not receive.
@dataclasses.dataclass
class RegressionState:
    y: torch.Tensor


# A synthetic regression task: the targets are a fixed random linear function of the inputs plus noise.
class RegressionDataset(Dataset):
    def __init__(self, num_samples: int, num_features: int):
        generator = torch.Generator().manual_seed(0)
        self._x = torch.randn(num_samples, num_features, generator=generator)
        true_weight = torch.randn(num_features, 1, generator=generator)
        self._y = self._x @ true_weight + 0.1 * torch.randn(num_samples, 1, generator=generator)

    def __getitem__(self, index: int) -> RegressionSample:
        return RegressionSample(x=self._x[index], y=self._y[index])

    def __len__(self) -> int:
        return len(self._x)

    @staticmethod
    def collate(samples: Sequence[RegressionSample]) -> RegressionBatch:
        return RegressionBatch(x=torch.stack([s.x for s in samples]), y=torch.stack([s.y for s in samples]))


# A plain PyTorch model. d9d needs two things from it: reset_parameters() for initialization on the
# device, and forward(inputs, shared), which is the signature of a pipeline stage.
class MLP(nn.Module):
    def __init__(self, num_features: int, hidden_size: int):
        super().__init__()
        self.up = nn.Linear(num_features, hidden_size)
        self.down = nn.Linear(hidden_size, 1)

    def forward(self, inputs: RegressionInput, shared: None) -> torch.Tensor:
        return self.down(torch.relu(self.up(inputs.x)))

    def reset_parameters(self):
        self.up.reset_parameters()
        self.down.reset_parameters()


class MLPProvider(ModelProvider[MLP]):
    def initialize_model_stage(self, context: InitializeModelStageContext) -> InitializeModelStageResult[MLP]:
        model = MLP(num_features=16, hidden_size=64)
        return InitializeModelStageResult(model=model, state_mapper=identity_mapper_from_module(model))

    def parallelize_model_stage(self, context: ParallelizeModelStageContext[MLP]):
        # HSDP: shards along dp_cp_shard and replicates along the other dimensions. Size-1 dimensions are skipped.
        mesh = context.dist_context.mesh_for(DENSE_DOMAIN)
        parallelize_hsdp(context.model, mesh=mesh["dp_replicate", "dp_cp_shard", "cp_replicate"])

    def prepare_export_model_stage(self, context: PrepareExportModelStageContext[MLP]) -> PrepareExportModelStageResult:
        return PrepareExportModelStageResult(state_mapper=identity_mapper_from_module(context.model))


class RegressionTask(TrainTask[RegressionBatch, RegressionInput, None, torch.Tensor, RegressionState]):
    def build_forward_inputs(
        self, ctx: BuildForwardInputsContext[RegressionBatch]
    ) -> BuildForwardInputsResult[RegressionInput, None, RegressionState]:
        # The model gets the inputs. The targets wait in the state until compute_loss().
        return BuildForwardInputsResult(
            input=RegressionInput(x=ctx.batch.x), shared=None, state=RegressionState(y=ctx.batch.y)
        )

    def compute_loss(self, ctx: ComputeLossContext[torch.Tensor, RegressionState]) -> ComputeLossResult:
        loss = nn.functional.mse_loss(ctx.pipeline_results, ctx.state.y)
        return ComputeLossResult(loss=loss, loss_weight=None)


def main():
    config = ProjectConfig.model_validate(yaml.safe_load(Path(sys.argv[1]).read_text(encoding="utf-8")))

    trainer = TrainingConfigurator(
        mesh=config.mesh,
        parameters=config.trainer,
        task_provider=lambda ctx: RegressionTask(),
        model_provider=MLPProvider(),
        data_provider=AutoDataProvider(
            dataset_factory=lambda dist_context: RegressionDataset(num_samples=65536, num_features=16),
            collator=RegressionDataset.collate,
            config=config.data,
        ),
        optimizer_provider=AutoOptimizerProvider(config.optimizer),
        lr_scheduler_provider=AutoLRSchedulerProvider(config.lr_scheduler),
    ).configure()

    trainer.train()
    trainer.export(config.export_to, load_checkpoint=False)


if __name__ == "__main__":
    main()

The script has these parts:

  • ProjectConfig: One Pydantic model for the whole YAML file. Pydantic validates the file before the job starts. See Configuration.
  • RegressionSample, RegressionBatch, RegressionInput and RegressionState: Dataclasses for one sample of the dataset, the microbatch that collate() builds, what the model receives and what compute_loss() needs besides the model output.
  • RegressionDataset: A standard PyTorch Dataset. AutoDataProvider shards it across data-parallel ranks, batches it with the collator and groups the microbatches into steps. See Data Loading.
  • MLP: A standard nn.Module. It implements reset_parameters(), because d9d builds the model on the meta device and initializes it on the GPU. Its forward() takes inputs and shared, like every pipeline stage. See Model Design.
  • MLPProvider: Builds the model, distributes it and returns the state mappers for loading and export. parallelize_hsdp() shards the model along dp_cp_shard and replicates it along the other dimensions of the mesh slice. A job on one GPU does not call parallelize_model_stage(). See Horizontal Parallelism and Model Definition.
  • RegressionTask: Splits each microbatch into the model input and the state, and computes the loss. The five type parameters of TrainTask serve only type checking. See User Tasks.
  • main(): TrainingConfigurator wires the parts together. configure() builds the job, train() runs it and export() saves the weights. load_checkpoint=False exports the weights that are in memory after training, without loading the latest checkpoint first. See Training Loop.

The Configuration

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
# The parallelism degrees. All of them default to 1, so the job runs on one GPU.
mesh: {}

data:
  global_batch_size: 256  # Samples per step.
  microbatch_size: 64  # Samples per forward and backward pass.
  num_workers: 2
  shuffle: true

trainer:
  run:
    name: quickstart  # Also the name of the checkpoint directory.
  # schedule.total_steps is not set, so the job makes one pass over the data.
  logging:
    period_steps: 10  # How often the task metrics are logged. The loss is logged every step.
    tracker:
      provider: log  # Writes the logged values to the log.
  determinism:
    base_seed: 42
  gc:
    period_steps: 100
  checkpointing:
    save_dir: runs/checkpoints
    period_steps: 100
    num_to_keep: 2
  gradient_clipping:
    max_norm: 1.0

optimizer:
  name: adamw
  lr: 0.01

# A linear warm-up over the first 10% of the steps, then a cosine decay to 10% of the peak.
lr_scheduler:
  name: piecewise
  scheduler:
    initial_multiplier: 0.0
    phases:
      - mode: percentage
        percentage: 0.1
        target_multiplier: 1.0
        curve:
          type: linear
      - mode: rest
        target_multiplier: 0.1
        curve:
          type: cosine

export_to: runs/export

The job makes one pass over the 65,536 samples: 256 steps of 256 samples.

trainer holds only the required fields. The other fields keep their defaults. For a first job of your own, you change these:

  • run.name: The name of the run and of its checkpoint directory.
  • schedule.total_steps: The number of steps. Without it, the job makes one pass over the data.
  • checkpointing: Where and how often the job saves checkpoints.
  • logging.tracker: Where the logged values go. log writes them to the log, and null drops them. Other providers send them to W&B, TensorBoard, MLflow or Aim and need an extra. See Experiment Tracking.

Configuration describes every field and its default.

Next Steps