User Tasks

@dataclasses.dataclass
class MyState:
    target: torch.Tensor
    num_targets: torch.Tensor

# In build_forward_inputs:
target = ctx.batch.target
state = MyState(target=target, num_targets=(target != LM_IGNORE_INDEX).sum())
return BuildForwardInputsResult(input=..., shared=..., state=state)

# Later, in compute_loss (ctx.state is typed as MyState):
loss = loss_fn(ctx.pipeline_results, ctx.state.target)

# And in update_metrics:
ctx.metrics["num_targets"].update(ctx.state.num_targets)

import dataclasses

import torch

from d9d.core.dist_context import DistributedContext
from d9d.core.types import ScalarTree
from d9d.loop.control import (
    BuildForwardInputsContext,
    BuildForwardInputsResult,
    ComputeLossContext,
    ComputeLossResult,
    TrainTask,
)
from d9d.module.block.head import LM_IGNORE_INDEX, SequenceCausalLMHeadShared, SequenceCausalLMOutput
from d9d.module.model.io import SequenceHeadShared, SequenceInput, SequenceShared


@dataclasses.dataclass
class SFTBatch:
    input_ids: torch.Tensor
    position_ids: torch.Tensor
    labels: torch.Tensor


@dataclasses.dataclass
class SFTState:
    labels: torch.Tensor


class SFTTask(
    TrainTask[
        SFTBatch,
        SequenceInput,
        SequenceHeadShared[SequenceCausalLMHeadShared],
        SequenceCausalLMOutput,
        SFTState,
    ]
):
    def __init__(self, dist_ctx: DistributedContext):
        self._dist_ctx = dist_ctx

    def build_forward_inputs(
        self, ctx: BuildForwardInputsContext[SFTBatch]
    ) -> BuildForwardInputsResult[SequenceInput, SequenceHeadShared[SequenceCausalLMHeadShared], SFTState]:
        # ctx.batch contains the output of the collator.
        # The SharedInput routes position IDs to the backbone and labels to the head.
        return BuildForwardInputsResult(
            input=SequenceInput(input_ids=ctx.batch.input_ids),
            shared=SequenceHeadShared(
                sequence=SequenceShared(position_ids=ctx.batch.position_ids),
                head=SequenceCausalLMHeadShared(labels=ctx.batch.labels),
            ),
            state=SFTState(labels=ctx.batch.labels),
        )

    def dump_hparams(self) -> ScalarTree:
        return super().dump_hparams()

    def compute_loss(self, ctx: ComputeLossContext[SequenceCausalLMOutput, SFTState]) -> ComputeLossResult:
        cross_entropy = ctx.pipeline_results.cross_entropy

        # Count the valid tokens (labels that are not padding). The loss of a variable-length batch needs it.
        num_loss_tokens = (ctx.state.labels != LM_IGNORE_INDEX).sum()

        # Average loss per valid token.
        total_loss = cross_entropy.sum() / num_loss_tokens

        return ComputeLossResult(
            loss=total_loss,
            # loss_weight weights this microbatch in the gradient average across microbatches and ranks.
            # Weighting by token count gives the true per-token average when token counts differ.
            loss_weight=num_loss_tokens / 1000,
        )