Training Loop

Overview

The d9d.loop package provides the execution engine for distributed training.

The d9d Trainer separates the definition of the job (Models, Tasks, Data) from the execution of the job (Synchronization, Checkpointing, Profiling).

This allows the same code to run on a single GPU or a 1000-GPU Pipeline Parallel cluster without modifications.

Example

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

from d9d.loop.run import TrainingConfigurator

# Configure
trainer = TrainingConfigurator(
    mesh=mesh_params,                  # Physical cluster layout
    parameters=config,                 # Logic configuration (batch size, etc)

    # --- User Logic ---
    model_provider=...,                # How to build the model
    task_provider=...,                 # How to compute loss
    data_provider=...,                 # How to load data
    optimizer_provider=...,            # How to optimize
    lr_scheduler_provider=...          # LR scheduler
).configure()

# Execute
trainer.train()

Configuration & Construction

To ensure reproducibility, the Trainer is not instantiated directly with loose objects. It is built using the TrainingConfigurator and the dependency injection pattern.

This class binds the

• Infrastructure Configuration,
• Job Configuration,
• and User Logic (Providers)

into a Trainer object with prepared TrainJobState.

The Configuration Lifecycle

The TrainingConfigurator.configure() method does:

1. Distributed Context Initialization:

   • Constructs the global DistributedContext, therefore initializing all the required NCCL process groups and DeviceMeshes.

2. Seeding:

   • Sets distributed seeds using the configured base_seed. This ensures model initialization and other initial states are deterministic. More info.

3. Event Bus Initialization:

   • Creates the global EventBus. Tasks and providers use this to register custom hooks.
   • Triggers EVENT_TRAIN_CONFIG_STARTED event.

4. Task Instantiation:

   • Instantiates the TrainTask object using specified TrainTaskProvider.

5. Data Loader Construction:

   • Calls the DatasetProvider to get the dataset and wraps it into a DataLoader.
   • The DataLoader will move all the Tensor data to this worker's device automatically.
   • Triggers EVENT_TRAIN_DATA_LOADER_READY event.

6. Model Materialization:

   • The ModelStageFactory runs. This is the heavy lifting of initialization:
     1. Meta Init: ModelProvider creates the model on the meta device (no memory usage).
     2. Parallelization: ModelProvider applies DTensor sharding/replication to parameters.
     3. Materialization: Empty tensors are allocated on the actual GPU.
     4. Wait: Hard barrier to ensure all ranks allocated memory successfully.
     5. Parameter Reset: model.reset_parameters() is called to generate random weights on GPU.
     6. Source Loading (Optional): If configured, a pretrained checkpoint (e.g., from HF) is streamed into the model using ModelStateMapper.

7. Triggers EVENT_TRAIN_MODEL_STAGES_READY event.

8. Optimizer and LR Scheduler Setup:

   • OptimizerFactory iterates over the model parameters.
   • Calls OptimizerProvider and LRSchedulerProvider.
   • Triggers both EVENT_TRAIN_OPTIMIZER_READY and EVENT_TRAIN_LR_SCHEDULER_READY events.

9. State Assembly:

   • All components (including internal ones) are packed into the TrainJobState.
   • The Trainer is instantiated with this state and returned.

Execution

To run a train job, just call the .train() method on a Trainer object that is returned by configuration process.

The Training Lifecycle

The Trainer.train() method orchestrates the following lifecycle. It is critical to understand this flow when debugging distributed issues or checking for side effects.

1. Initialization & Recovery

Before the loop starts:

1. Global Synchronization: The trainer waits for all ranks to come online (barrier).
2. State Loading: The StateCheckpointer checks the filesystem.
   • If a checkpoint exists, it loads it into all the Stateful objects inside its _state.
   • If no checkpoint exists, it starts from the first step.
3. Context Entry: The trainer enters several context managers:
   • UI: Renders a progress bar.
   • Logging: Initiates a new run in selected experiment tracker and dumps run hyperparameters there. More info.
   • Garbage Collector: Disables automatic Python garbage collection.
   • Profiler: Starts torch.profiler hooks. More info.
   • Gradient Manager: Sets up backward hooks for synchronizing gradient states by all-reduce.
   • Gradient Clipper: Looks for model parameters which gradients will be registered for clipping.
4. Ready Hook Trigger: EVENT_TRAIN_READY is fired to mark the start of the primary train sequence.

2. The Step Loop

For every global step (step), the trainer performs the following actions in strict order:

1. Triggers EVENT_TRAIN_STEP_PRE event.

2. Microbatch Execution

   • Triggers EVENT_TRAIN_FORWARD_BACKWARD_PRE event.
   • The DataLoader yields a "Batch Group" containing \(N\) microbatches (calculated automatically based on BatchingConfig).
   • We delegate to the TrainTask for mapping data before feeding it into the model.
   • The gradients will be accumulated locally using either regular multiple forward-backward calls if pipeline parallelism is disabled, either using our internal pipelining API. We delegate to TrainTask to compute loss values between forward and backward passes.
   • Last gradient accumulation triggers all-reduce synchronization. Communications may start overlapping here.
   • We delegate to TrainTask to accumulate local metrics (e.g., token counts, accuracy) into the Metric state.
   • Triggers EVENT_TRAIN_FORWARD_BACKWARD_POST event.

3. Metric Synchronization

   • Metric Sync Trigger: JobLogger triggers an async reduction of all metrics across the world. More info.

4. Gradient Synchronization

   • Wait & Scale: The GradientManager waits for all backward hooks to finish. It synchronizes the total weighted loss across the world to determine the scaling factor, then divides all gradients by this factor (essential for correct averaging when batch sizes vary due to masking/packing). More info.

5. Gradient Clipping

   • The GradientClipper calculates the global L2 norm of all parameters.
   • If max_norm is set, gradients are modified in-place.
   • The total norm is logged.
   • More info.

6. Optimization

   • Triggers EVENT_TRAIN_OPTIMIZER_STEP_PRE event.
   • Step: The Optimizer updates model parameters.
   • Schedule: The LRScheduler updates the learning rate for the next step.
   • Zero Grad: The GradientManager clears gradients for the next iteration.
   • Triggers EVENT_TRAIN_OPTIMIZER_STEP_POST event.

7. Logging & Maintenance

   • Log: Metrics are finalized and written to the tracker.
   • GC: ManualGarbageCollector runs if the current step matches the GC period.
   • Event-Based Logic: Triggers EVENT_TRAIN_STEP_POST event.
   • Advance: The Stepper increments the step count.

8. Checkpointing

   • If the current step matches checkpointing.period_steps, checkpointing is triggered. This acts as a global barrier.

3. Finalization

1. Event-specific: The system triggers EVENT_TRAIN_FINISHED event.
2. Task-specific: We delegate to the TrainTask to do its specific finalization work.

Sleep & Wake (Colocated RL)

In colocated RL, a rollout/inference engine shares the same physical GPUs as the Trainer. The two cannot occupy the device at once, so the Trainer must temporarily hand the GPU back. The Trainer.sleep() / Trainer.wake() pair does exactly this: sleep() releases the training state to host memory and frees the device cache; wake() restores it. The underlying mechanism is the State Offloading subsystem.

1
2
3
4
5
6

trainer = TrainingConfigurator(...).configure()

# ... after a training step, between steps ...
trainer.sleep()                 # offload training state, free the GPU
rollout_engine.generate(...)    # colocated inference runs on the freed device
trainer.wake()                  # restore training state, resume stepping

What gets offloaded

A sleep() selects subsystems via SleepTag. The default, SleepTag.TENSOR_STATES, offloads all GPU tensor state as a single unit:

• Model parameters and buffers (TrackedModules),
• Optimizer state (PipelinedOptimizer),
• Gradient buckets and the residual loss accumulator (GradientManager).

SleepTag.COMMS (NCCL process groups) is reserved but not yet implemented - requesting it raises NotImplementedError.

The round trip preserves object identity, optimizer state-dict keys and DTensor metadata, so any external reference (e.g. a frozen reference model held by a task) stays valid after waking. See the round-trip invariant for details.

Orchestration & ordering

sleep() performs, in order:

1. Fire EVENT_TRAIN_SLEEP_PRE (state still fully GPU-resident) and barrier.
2. Offload gradient state → optimizer state → model parameters and buffers.
3. Barrier and fire EVENT_TRAIN_SLEEP_POST.

wake() restores in the reverse order (model → optimizer → gradient state), bracketed by EVENT_TRAIN_WAKE_PRE / EVENT_TRAIN_WAKE_POST and barriers.