Write Your Own Job

# One microbatch of token sequences, padded to one length.
@dataclasses.dataclass
class TextBatch:
    input_ids: torch.Tensor
    labels: torch.Tensor
    position_ids: torch.Tensor


class TextDataset(Dataset):
    ...

    @staticmethod
    def collate(samples: Sequence[torch.Tensor]) -> TextBatch:
        # d9d models do not shift the labels, so shift them here.
        return TextBatch(
            input_ids=pad_stack_1d([tokens[:-1] for tokens in samples], pad_value=0),
            labels=pad_stack_1d([tokens[1:] for tokens in samples], pad_value=LM_IGNORE_INDEX),
            position_ids=pad_stack_1d([torch.arange(len(tokens) - 1) for tokens in samples], pad_value=0),
        )


def build_dataset(config: DataConfig, dist_context: DistributedContext) -> Dataset:
    # Rank 0 downloads the dataset first. The other ranks then read it from the cache.
    with dist_context.main_process_first():
        data = datasets.load_dataset(config.dataset, config.dataset_config, split=config.split)
    texts = [text for text in data[config.text_column] if text.strip()]
    return TextDataset(texts, Tokenizer.from_file(str(config.tokenizer)), config.max_length)


# The data provider of TrainingConfigurator. It calls build_dataset() once on every rank
# and TextDataset.collate() for every microbatch.
data_provider = AutoDataProvider(
    dataset_factory=lambda dist_context: build_dataset(config.data, dist_context),
    collator=TextDataset.collate,
    config=config.auto_data,
)

class Block(nn.Module):
    def __init__(self, hidden_size: int):
        super().__init__()
        self.norm = nn.LayerNorm(hidden_size)
        self.linear = nn.Linear(hidden_size, hidden_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.linear(self.norm(x))

    def reset_parameters(self):
        self.norm.reset_parameters()
        self.linear.reset_parameters()


class ResidualMLP(nn.Module):
    def __init__(self, num_features: int, hidden_size: int, num_blocks: int):
        super().__init__()
        self.up = nn.Linear(num_features, hidden_size)
        self.blocks = nn.ModuleList(Block(hidden_size) for _ in range(num_blocks))
        self.down = nn.Linear(hidden_size, 1)
        self.register_buffer("input_scale", torch.empty(num_features))

    def forward(self, inputs: RegressionInput, shared: None) -> torch.Tensor:
        x = self.up(inputs.x * self.input_scale)
        for block in self.blocks:
            x = block(x)
        return self.down(x)

    def reset_parameters(self):
        self.up.reset_parameters()
        for block in self.blocks:
            block.reset_parameters()
        self.down.reset_parameters()
        self.input_scale.fill_(1.0)

class ModelConfig(BaseModel):
    num_features: int
    hidden_size: int
    num_blocks: int


class ResidualMLPProvider(ModelProvider[ResidualMLP]):
    def __init__(self, config: ModelConfig):
        self._config = config

    # Builds the model and returns the state mapper for loading.
    def initialize_model_stage(self, context: InitializeModelStageContext) -> InitializeModelStageResult[ResidualMLP]:
        model = ResidualMLP(
            num_features=self._config.num_features,
            hidden_size=self._config.hidden_size,
            num_blocks=self._config.num_blocks,
        )
        # The model trains from scratch and keeps its own parameter names, so the mapper changes nothing.
        return InitializeModelStageResult(model=model, state_mapper=identity_mapper_from_module(model))

    # Distributes the model. The loop calls it only in a distributed run, not on one GPU.
    def parallelize_model_stage(self, context: ParallelizeModelStageContext[ResidualMLP]):
        mesh = context.dist_context.mesh_for(DENSE_DOMAIN)
        parallelize_hsdp(context.model, mesh=mesh["dp_replicate", "dp_cp_shard", "cp_replicate"])

    # Returns the state mapper for export.
    def prepare_export_model_stage(
        self, context: PrepareExportModelStageContext[ResidualMLP]
    ) -> PrepareExportModelStageResult:
        return PrepareExportModelStageResult(state_mapper=identity_mapper_from_module(context.model))


class ProjectConfig(BaseModel):
    model: ModelConfig
    ...


# The model provider of TrainingConfigurator in main().
model_provider = ResidualMLPProvider(config.model)

# The type parameters describe the microbatch, input, shared, the model output and state. They serve only
# type checking. Each part can be a dataclass, a dict, or None if the job has no such data.
class RegressionTask(TrainTask[RegressionBatch, RegressionInput, None, torch.Tensor, RegressionState]):
    # Gets the microbatch that collate() built and splits it into three parts.
    def build_forward_inputs(
        self, ctx: BuildForwardInputsContext[RegressionBatch]
    ) -> BuildForwardInputsResult[RegressionInput, None, RegressionState]:
        return BuildForwardInputsResult(
            # The model receives it as `inputs`.
            input=RegressionInput(x=ctx.batch.x),
            # The model receives it as `shared`. With pipeline parallelism, every stage receives it.
            shared=None,
            # What compute_loss() and update_metrics() need but the model does not return.
            state=RegressionState(y=ctx.batch.y),
        )

    # Gets the model output and the state of the same microbatch.
    def compute_loss(self, ctx: ComputeLossContext[torch.Tensor, RegressionState]) -> ComputeLossResult:
        loss = nn.functional.mse_loss(ctx.pipeline_results, ctx.state.y)
        # The loop averages the losses of all microbatches and ranks, weighted by loss_weight. None means 1.
        return ComputeLossResult(loss=loss, loss_weight=None)


# The task provider of TrainingConfigurator in main(). The loop calls it after it builds the distributed
# context and passes the context in ctx.dist_context.
task_provider = lambda ctx: RegressionTask()

from d9d.metric.impl.aggregation import SumMetric


class RegressionTask(TrainTask[...]):
    ...

    # Called once, when the loop builds the job. The keys are the names in the log.
    def create_metrics(self, ctx: CreateMetricsContext) -> CreateMetricsResult:
        return CreateMetricsResult(metrics={"num_samples": SumMetric()})

    # Called after compute_loss() for each microbatch. It gets the state, not the model output, so
    # build_forward_inputs() must put everything that a metric needs into the state.
    def update_metrics(self, ctx: UpdateMetricsContext[RegressionState]):
        ctx.metrics["num_samples"].update(torch.ones_like(ctx.state.y))

model:
  num_features: 16
  hidden_size: 64
  num_blocks: 2