Skip to content

Model Definition

About

The ModelProvider controls the lifecycle of the nn.Module. In distributed training, a model is not just created. It must be initialized, parallelized and mapped for loading from a checkpoint.

How to Write a ModelProvider

Choose a Model

Choose a model from the d9d catalogue, or create your own.

Implement initialize_model_stage(...)

This method builds the nn.Module for one pipeline parallel stage, in the target torch.dtype.

The loop calls it on the meta device, so you must not load model weights here. Instead, return a state mapper that maps the weights on disk to the weights in memory.

You can also apply PEFT methods and other architecture patches here. The returned state mapper must reflect the changes they make.

Implement parallelize_model_stage(...)

This method applies a horizontal parallelism strategy to the model in place. The loop calls it only for distributed runs.

For d9d models, you can use the default strategies. For example, apply parallelize_qwen3_moe_model to the backbone and the routine for each head type, such as parallelize_causal_lm_head (reference).

For a custom model, see the horizontal parallelism docs and the reference implementations.

Implement prepare_export_model_stage(...) for Export

You must implement this method if the job calls Trainer.export(). Other jobs can skip it. Without it, export() raises a NotImplementedError.

This method returns a state mapper for the final export. It converts the in-memory model state to the state saved on disk.

It usually reverses the state mapper returned by initialize_model_stage(...).

Usage

 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
from pydantic import BaseModel

from d9d.core.types import ScalarTree
from d9d.loop.control import (
    InitializeModelStageContext,
    InitializeModelStageResult,
    ModelProvider,
    ParallelizeModelStageContext,
    PrepareExportModelStageContext,
    PrepareExportModelStageResult,
)
from d9d.model_state.mapper.adapters import identity_mapper_from_module
from d9d.module.block.hidden_states_aggregator import HiddenStatesAggregationMode
from d9d.module.model import DecoderForCausalLM
from d9d.module.model.qwen3_moe import Qwen3MoEModel, Qwen3MoEParameters
from d9d.module.parallelism.model import parallelize_causal_lm_head, parallelize_qwen3_moe_model


class ModelProviderConfig(BaseModel):
    model: Qwen3MoEParameters  # Hyperparameters for the Qwen3 MoE backbone
    checkpointing: bool  # Enable activation checkpointing to save GPU memory


class ProjectModelProvider(ModelProvider[DecoderForCausalLM[Qwen3MoEModel]]):
    def __init__(self, config: ModelProviderConfig):
        self._config = config

    def initialize_model_stage(self, context: InitializeModelStageContext) -> InitializeModelStageResult:
        # Compose the Qwen3 MoE backbone with a single causal LM head and cast it to bf16.
        backbone = Qwen3MoEModel(
            params=self._config.model,
            stage=context.stage,
            hidden_states_snapshot_mode=HiddenStatesAggregationMode.no,
            enable_checkpointing=self._config.checkpointing,
        )
        model = DecoderForCausalLM(backbone, context.stage).bfloat16()

        return InitializeModelStageResult(
            model=model,
            state_mapper=identity_mapper_from_module(model),
        )

    def parallelize_model_stage(self, context: ParallelizeModelStageContext):
        # Apply the Qwen3 MoE parallelism routine to the backbone and the head routine to the head.
        # You can apply your own horizontal parallelism strategy here.
        parallelize_qwen3_moe_model(context.dist_context, context.model.model, context.stage)
        if context.stage.is_current_stage_last:
            parallelize_causal_lm_head(context.model.head, context.dist_context)

    def prepare_export_model_stage(self, context: PrepareExportModelStageContext) -> PrepareExportModelStageResult:
        # Export the model weights as they are.
        return PrepareExportModelStageResult(
            state_mapper=identity_mapper_from_module(context.model),
        )

    def dump_hparams(self) -> ScalarTree:
        return self._config.model_dump(mode="json")

API Reference

d9d.loop.control.model_provider

Interface for building, parallelizing and exporting model stages.

InitializeModelStageContext dataclass

Context data required for initializing a specific model pipeline stage.

Attributes:

Name Type Description
dist_context DistributedContext

The distributed execution context.
stage PipelineStageInfo

Metadata describing the current pipeline stage being initialized.

InitializeModelStageResult dataclass

Bases: Generic[TModel]

The result of initializing a model stage.

Attributes:

Name Type Description
model TModel

The PyTorch module.
state_mapper ModelStateMapper

The mapper defining how to load weights into this module.

ModelProvider

Bases: ABC, Generic[TModel]

Abstract base class for the lifecycle of a distributed model.

The provider initializes, parallelizes (shards, replicates, etc.) and prepares the model for export.

dump_hparams()

Exports hyperparameters associated with this model for logging.

Returns:

Type Description
ScalarTree

A dictionary of hyperparameter names and values.
initialize_model_stage(context) abstractmethod

Initializes the model architecture for a specific pipeline stage.

It constructs the nn.Module for the requested stage.

Construction runs on the meta device, so the method must not load weights. Instead, it returns a ModelStateMapper that maps checkpoint weights to the parameters of the new module.

The method can change the architecture, e.g. inject LoRA adapters, if the returned mapper reflects the new structure.

Parameters:

Name Type Description Default
context InitializeModelStageContext

Context for this operation.

 required

Returns:

Type Description
InitializeModelStageResult[TModel]

The model stage and its state mapper.
parallelize_model_stage(context) abstractmethod

Converts the model parameters into distributed tensors (DTensor).

Implementations must modify the model in place. They replicate or shard each parameter according to the chosen parallelism strategies.

Parameters:

Name Type Description Default
context ParallelizeModelStageContext[TModel]

Context for this operation.

 required
prepare_export_model_stage(context)

Prepares the state mapper required for saving the model to disk.

The mapper defines how the in-memory model structure maps back to the checkpoint format. A provider must implement this method if its job calls Trainer.export().

Parameters:

Name Type Description Default
context PrepareExportModelStageContext[TModel]

Context for this operation.

 required

Returns:

Type Description
PrepareExportModelStageResult

The state mapper for export.

Raises:

Type Description
NotImplementedError

If the provider does not implement this method.
register_events(context)

Registers model-specific event subscriptions.

Parameters:

Name Type Description Default
context RegisterModelEventsContext

Context with the distributed context and the event bus.

 required

ParallelizeModelStageContext dataclass

Bases: Generic[TModel]

Context data required for horizontally parallelizing a model stage.

Attributes:

Name Type Description
dist_context DistributedContext

The distributed execution context.
stage PipelineStageInfo

Metadata describing the current pipeline stage.
model TModel

The PyTorch module to be parallelized.

PrepareExportModelStageContext dataclass

Bases: Generic[TModel]

Context data required for preparing a model stage for export.

Attributes:

Name Type Description
dist_context DistributedContext

The distributed execution context.
model TModel

The PyTorch module to be exported.

PrepareExportModelStageResult dataclass

The result of preparing a model stage for export.

Attributes:

Name Type Description
state_mapper ModelStateMapper

The mapper defining how model parameters map to disk storage.

RegisterModelEventsContext dataclass

Context for registering model-specific events.

Attributes:

Name Type Description
dist_context DistributedContext

The distributed execution context.
event_bus EventBus

The event bus for subscribing to events.