Fine-Tune a Hugging Face Model

About

This guide loads Qwen3-0.6B from its Hugging Face checkpoint, fine-tunes it on WikiText-2 and exports it back in the Hugging Face format. A second config trains LoRA adapters instead of all weights. Another script measures the perplexity of a model with the inference loop. The code is in example/guide/finetune_qwen3.

Hugging Face integration will change

d9d will integrate with Hugging Face more closely (#69), so this guide will change. Today, you copy the model parameters from config.json by hand and copy the config and the tokenizer files after export.

Run the Example

Install d9d and get the repository as in the Quickstart. The example also reads the data with the datasets and tokenizers packages:

pip install datasets tokenizers
cd d9d/example/guide/finetune_qwen3
hf download Qwen/Qwen3-0.6B --local-dir models/Qwen3-0.6B
python finetune.py config.yaml

The job trains for 100 steps on one GPU. It saves a checkpoint to runs/checkpoints/finetune-qwen3/ and the exported model to runs/export/.

Model Parameters

The model section of the config holds the Qwen3DenseParameters. Copy them from the config.json file of the Hugging Face model:

Hugging Face `config.json`	d9d `model`
`hidden_size`, `intermediate_size`, `num_attention_heads`, `num_key_value_heads`, `head_dim`, `rms_norm_eps`	`layer.<same name>`
`num_hidden_layers`	`num_hidden_layers`
`rope_theta`	`rope_base`
`max_position_embeddings`	`max_position_ids`
`vocab_size`	`split_vocab_size.regular`, with `split_vocab_order` set to `["regular"]`

For Qwen3-0.6B, the section is:

model:
  layer:
    hidden_size: 1024
    intermediate_size: 3072
    num_attention_heads: 16
    num_key_value_heads: 8
    rms_norm_eps: 1.0e-06
    head_dim: 128
  num_hidden_layers: 28
  rope_base: 1000000
  max_position_ids: 40960
  split_vocab_size:
    regular: 151936
  split_vocab_order:
    - regular

The Hugging Face mappers support only one vocabulary split.

Loading the Hugging Face Model

trainer.model_stage_factory.source_checkpoint points to the downloaded model directory:

trainer:
  model_stage_factory:
    source_checkpoint: models/Qwen3-0.6B

The job streams the weights into the model through the state mapper that initialize_model_stage() returns (see The Model Build). The model provider also handles LoRA and export, which the sections below describe:

class Qwen3Provider(ModelProvider[DecoderForCausalLM[Qwen3DenseModel]]):
    def __init__(self, params: Qwen3DenseParameters, lora: LoRAConfig | None):
        self._params = params
        self._lora = LoRA(lora) if lora is not None else None

    def initialize_model_stage(self, context: InitializeModelStageContext) -> InitializeModelStageResult:
        backbone = Qwen3DenseModel(
            params=self._params,
            stage=context.stage,
            hidden_states_snapshot_mode=HiddenStatesAggregationMode.no,
            enable_checkpointing=True,
        )
        model = DecoderForCausalLM(backbone, context.stage).bfloat16()

        # Maps the Hugging Face checkpoint keys to the d9d model keys.
        state_mapper = mapper_from_huggingface_qwen3_dense_for_causal_lm(self._params)

        if self._lora is not None:
            # Wraps the matching layers in LoRA adapters and freezes everything else.
            # The LoRA mapper then moves the original weights of the wrapped layers to their new keys.
            lora_mapper = inject_peft_and_freeze(self._lora, model)
            state_mapper = ModelStateMapperSequential([state_mapper, lora_mapper])

        return InitializeModelStageResult(model=model, state_mapper=state_mapper)

    def parallelize_model_stage(self, context: ParallelizeModelStageContext):
        parallelize_qwen3_dense_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:
        if self._lora is not None:
            # Merges the adapters into the base weights, so the export has the original architecture.
            merge_peft(self._lora, context.model)

        # Exports the model in the Hugging Face format.
        return PrepareExportModelStageResult(state_mapper=mapper_to_huggingface_qwen3_dense_for_causal_lm(self._params))

See Model State Mapper.

Optimizer

The provider casts the model to bf16 with .bfloat16(). The config keeps the gradients in fp32 and uses stochastic_adamw, which updates bf16 weights with stochastic rounding (see Stochastic Optimizers):

trainer:
  gradient_manager:
    grad_dtype: float32

optimizer:
  name: stochastic_adamw
  lr: 1.0e-05
  state_dtype: float32

The adamw optimizer requires the gradients to have the dtype of the weights, so it fails with this config.

Loss

The causal LM head computes the loss itself, so the task passes the labels to it in shared and keeps only the number of label tokens in the state:

@dataclasses.dataclass
class CausalLMState:
    num_tokens: torch.Tensor


class CausalLMTask(
    TrainTask[
        TextBatch,
        SequenceInput,
        SequenceHeadShared[SequenceCausalLMHeadShared],
        SequenceCausalLMOutput,
        CausalLMState,
    ]
):
    def build_forward_inputs(
        self, ctx: BuildForwardInputsContext[TextBatch]
    ) -> BuildForwardInputsResult[SequenceInput, SequenceHeadShared[SequenceCausalLMHeadShared], CausalLMState]:
        labels = ctx.batch.labels
        return BuildForwardInputsResult(
            input=SequenceInput(input_ids=ctx.batch.input_ids),
            shared=SequenceHeadShared(
                sequence=SequenceShared(position_ids=ctx.batch.position_ids),
                head=SequenceCausalLMHeadShared(labels=labels),
            ),
            state=CausalLMState(num_tokens=(labels != LM_IGNORE_INDEX).sum()),
        )

    def compute_loss(self, ctx: ComputeLossContext[SequenceCausalLMOutput, CausalLMState]) -> ComputeLossResult:
        # The head returns the per-token cross-entropy loss. Padding positions hold zeros.
        num_tokens = ctx.state.num_tokens
        loss = ctx.pipeline_results.cross_entropy.sum() / num_tokens

        # Weighting each microbatch by its token count gives the mean loss per token across microbatches and ranks.
        return ComputeLossResult(loss=loss, loss_weight=num_tokens.float())

Returning the number of tokens as the loss_weight makes the loop average the loss per token across microbatches and ranks.

Export to Hugging Face

prepare_export_model_stage() of the provider returns mapper_to_huggingface_qwen3_dense_for_causal_lm, so Trainer.export() writes the weights with the Hugging Face keys. The export has only the .safetensors files. Copy the config and the tokenizer from the original model to load it with Transformers:

cp models/Qwen3-0.6B/{config.json,generation_config.json,tokenizer.json,tokenizer_config.json} runs/export/

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("runs/export")
model = AutoModelForCausalLM.from_pretrained("runs/export", dtype="bfloat16").cuda()

inputs = tokenizer("The history of the Roman Empire", return_tensors="pt").to("cuda")
print(tokenizer.decode(model.generate(**inputs, max_new_tokens=25)[0]))

Qwen3-0.6B ties its input embeddings to its LM head. d9d trains them as two separate weights, so Transformers logs that it does not tie them. Set "tie_word_embeddings": false in the copied config.json to silence the message.

Fine-Tune with LoRA

config_lora.yaml adds a lora section, which trains LoRA adapters on the attention projections:

lora:
  module_name_pattern: .*self_attn\.(q|k|v|o)_proj
  params:
    r: 16
    alpha: 32
    dropout: 0.0

1	`python finetune.py config_lora.yaml`

The provider changes in three places:

Injection: initialize_model_stage() wraps the matching layers and freezes all other weights. It chains the LoRA mapper after the Hugging Face mapper, because the wrapped layers move their weights to new keys.
Checkpoints: The config sets checkpoint_only_trainable_parameters, so the checkpoints hold only the adapters: 46 MiB instead of 7.1 GiB.
Export: prepare_export_model_stage() merges the adapters into the base weights first. The export in runs/export-lora/ has the original architecture. Copy the config and the tokenizer into it as above to load it with Transformers.

Evaluate the Model

evaluate.py computes the perplexity on the validation split with the inference loop. It takes the model state to evaluate as its second argument. It loads it through the same Hugging Face mapper, so it reads the original model and the exports alike:

python evaluate.py evaluate.yaml models/Qwen3-0.6B
python evaluate.py evaluate.yaml runs/export
python evaluate.py evaluate.yaml runs/export-lora

Each run logs the mean loss and the perplexity on the WikiText-2 validation split.

The script reuses the dataset and the model provider of finetune.py. Its task sums the loss and the number of tokens of each microbatch and computes the perplexity at the end:

@dataclasses.dataclass
class PerplexityState:
    num_tokens: torch.Tensor


class PerplexityTask(
    InferenceTask[
        TextBatch,
        SequenceInput,
        SequenceHeadShared[SequenceCausalLMHeadShared],
        SequenceCausalLMOutput,
        PerplexityState,
    ]
):
    def __init__(self, dist_context: DistributedContext):
        self._dist_context = dist_context
        self._loss_sum = torch.zeros((), dtype=torch.float64, device="cuda")
        self._num_tokens = torch.zeros((), dtype=torch.float64, device="cuda")

    def build_forward_inputs(
        self, ctx: BuildForwardInputsContext[TextBatch]
    ) -> BuildForwardInputsResult[SequenceInput, SequenceHeadShared[SequenceCausalLMHeadShared], PerplexityState]:
        labels = ctx.batch.labels
        return BuildForwardInputsResult(
            input=SequenceInput(input_ids=ctx.batch.input_ids),
            shared=SequenceHeadShared(
                sequence=SequenceShared(position_ids=ctx.batch.position_ids),
                head=SequenceCausalLMHeadShared(labels=labels),
            ),
            state=PerplexityState(num_tokens=(labels != LM_IGNORE_INDEX).sum()),
        )

    def process_outputs(self, ctx: ProcessOutputsContext[SequenceCausalLMOutput, PerplexityState]):
        # The head returns the per-token cross-entropy loss. Padding positions hold zeros.
        self._loss_sum += ctx.pipeline_results.cross_entropy.sum()
        self._num_tokens += ctx.state.num_tokens

    def finalize(self, ctx: FinalizeContext) -> None:
        # Ranks that host no last pipeline stage add zeros, so a sum over all ranks covers every batch once.
        if self._dist_context.mesh_params.is_distributed:
            group = self._dist_context.mesh_for(FLAT_DOMAIN).get_group()
            dist.all_reduce(self._loss_sum, group=group)
            dist.all_reduce(self._num_tokens, group=group)

        mean_loss = (self._loss_sum / self._num_tokens).item()
        if self._dist_context.is_main_process:
            self._dist_context.logger.info(f"Mean loss {mean_loss:.4f}, perplexity {math.exp(mean_loss):.2f}")

More GPUs

parallelize_model_stage() applies the default Qwen3 strategy: HSDP over the data-parallel dimensions. Set data_parallel_replicate or data_parallel_shard in mesh and launch with torchrun, as described in Running Jobs. See Qwen3 Dense for the supported parallelism.

The Code

finetune.pyevaluate.py

import dataclasses
import sys
from collections.abc import Sequence
from pathlib import Path

import datasets
import torch
import yaml
from d9d.core.dist_context import DeviceMeshParameters, DistributedContext
from d9d.dataset import pad_stack_1d
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.compose import ModelStateMapperSequential
from d9d.module.block.head import LM_IGNORE_INDEX, SequenceCausalLMHeadShared, SequenceCausalLMOutput
from d9d.module.block.hidden_states_aggregator import HiddenStatesAggregationMode
from d9d.module.model import DecoderForCausalLM
from d9d.module.model.io import SequenceHeadShared, SequenceInput, SequenceShared
from d9d.module.model.qwen3_dense import (
    Qwen3DenseModel,
    Qwen3DenseParameters,
    mapper_from_huggingface_qwen3_dense_for_causal_lm,
    mapper_to_huggingface_qwen3_dense_for_causal_lm,
)
from d9d.module.parallelism.model import parallelize_causal_lm_head, parallelize_qwen3_dense_model
from d9d.peft import inject_peft_and_freeze, merge_peft
from d9d.peft.lora import LoRA, LoRAConfig
from pydantic import BaseModel, ConfigDict
from tokenizers import Tokenizer
from torch.utils.data import Dataset


class DataConfig(BaseModel):
    model_config = ConfigDict(extra="forbid")

    dataset: str
    dataset_config: str | None
    split: str
    text_column: str
    tokenizer: Path
    max_length: int


class ProjectConfig(BaseModel):
    model_config = ConfigDict(extra="forbid")

    mesh: DeviceMeshParameters
    model: Qwen3DenseParameters
    lora: LoRAConfig | None
    data: DataConfig
    auto_data: AutoDataConfig
    trainer: TrainerConfig
    optimizer: AutoOptimizerConfig
    lr_scheduler: AutoLRSchedulerConfig
    export_to: Path


# 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):
    def __init__(self, texts: list[str], tokenizer: Tokenizer, max_length: int):
        self._samples = []
        for text in texts:
            # One more token than max_length, because the labels are the inputs shifted by one.
            ids = tokenizer.encode(text).ids[: max_length + 1]
            if len(ids) > 1:
                self._samples.append(torch.tensor(ids, dtype=torch.long))

    def __getitem__(self, index: int) -> torch.Tensor:
        return self._samples[index]

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

    @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)


class Qwen3Provider(ModelProvider[DecoderForCausalLM[Qwen3DenseModel]]):
    def __init__(self, params: Qwen3DenseParameters, lora: LoRAConfig | None):
        self._params = params
        self._lora = LoRA(lora) if lora is not None else None

    def initialize_model_stage(self, context: InitializeModelStageContext) -> InitializeModelStageResult:
        backbone = Qwen3DenseModel(
            params=self._params,
            stage=context.stage,
            hidden_states_snapshot_mode=HiddenStatesAggregationMode.no,
            enable_checkpointing=True,
        )
        model = DecoderForCausalLM(backbone, context.stage).bfloat16()

        # Maps the Hugging Face checkpoint keys to the d9d model keys.
        state_mapper = mapper_from_huggingface_qwen3_dense_for_causal_lm(self._params)

        if self._lora is not None:
            # Wraps the matching layers in LoRA adapters and freezes everything else.
            # The LoRA mapper then moves the original weights of the wrapped layers to their new keys.
            lora_mapper = inject_peft_and_freeze(self._lora, model)
            state_mapper = ModelStateMapperSequential([state_mapper, lora_mapper])

        return InitializeModelStageResult(model=model, state_mapper=state_mapper)

    def parallelize_model_stage(self, context: ParallelizeModelStageContext):
        parallelize_qwen3_dense_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:
        if self._lora is not None:
            # Merges the adapters into the base weights, so the export has the original architecture.
            merge_peft(self._lora, context.model)

        # Exports the model in the Hugging Face format.
        return PrepareExportModelStageResult(state_mapper=mapper_to_huggingface_qwen3_dense_for_causal_lm(self._params))




@dataclasses.dataclass
class CausalLMState:
    num_tokens: torch.Tensor


class CausalLMTask(
    TrainTask[
        TextBatch,
        SequenceInput,
        SequenceHeadShared[SequenceCausalLMHeadShared],
        SequenceCausalLMOutput,
        CausalLMState,
    ]
):
    def build_forward_inputs(
        self, ctx: BuildForwardInputsContext[TextBatch]
    ) -> BuildForwardInputsResult[SequenceInput, SequenceHeadShared[SequenceCausalLMHeadShared], CausalLMState]:
        labels = ctx.batch.labels
        return BuildForwardInputsResult(
            input=SequenceInput(input_ids=ctx.batch.input_ids),
            shared=SequenceHeadShared(
                sequence=SequenceShared(position_ids=ctx.batch.position_ids),
                head=SequenceCausalLMHeadShared(labels=labels),
            ),
            state=CausalLMState(num_tokens=(labels != LM_IGNORE_INDEX).sum()),
        )

    def compute_loss(self, ctx: ComputeLossContext[SequenceCausalLMOutput, CausalLMState]) -> ComputeLossResult:
        # The head returns the per-token cross-entropy loss. Padding positions hold zeros.
        num_tokens = ctx.state.num_tokens
        loss = ctx.pipeline_results.cross_entropy.sum() / num_tokens

        # Weighting each microbatch by its token count gives the mean loss per token across microbatches and ranks.
        return ComputeLossResult(loss=loss, loss_weight=num_tokens.float())




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: CausalLMTask(),
        model_provider=Qwen3Provider(config.model, config.lora),
        data_provider=AutoDataProvider(
            dataset_factory=lambda dist_context: build_dataset(config.data, dist_context),
            collator=TextDataset.collate,
            config=config.auto_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()

import dataclasses
import math
import sys
from pathlib import Path

import torch
import torch.distributed as dist
import yaml
from d9d.core.dist_context import FLAT_DOMAIN, DeviceMeshParameters, DistributedContext
from d9d.loop.auto import AutoDataConfig, AutoDataProvider
from d9d.loop.config import InferenceConfig
from d9d.loop.control import (
    BuildForwardInputsContext,
    BuildForwardInputsResult,
    FinalizeContext,
    InferenceTask,
    ProcessOutputsContext,
)
from d9d.loop.run import InferenceConfigurator
from d9d.module.block.head import LM_IGNORE_INDEX, SequenceCausalLMHeadShared, SequenceCausalLMOutput
from d9d.module.model.io import SequenceHeadShared, SequenceInput, SequenceShared
from d9d.module.model.qwen3_dense import Qwen3DenseParameters
from finetune import DataConfig, Qwen3Provider, TextBatch, TextDataset, build_dataset
from pydantic import BaseModel, ConfigDict


class EvaluationConfig(BaseModel):
    model_config = ConfigDict(extra="forbid")

    mesh: DeviceMeshParameters
    model: Qwen3DenseParameters
    data: DataConfig
    auto_data: AutoDataConfig
    inference: InferenceConfig


@dataclasses.dataclass
class PerplexityState:
    num_tokens: torch.Tensor


class PerplexityTask(
    InferenceTask[
        TextBatch,
        SequenceInput,
        SequenceHeadShared[SequenceCausalLMHeadShared],
        SequenceCausalLMOutput,
        PerplexityState,
    ]
):
    def __init__(self, dist_context: DistributedContext):
        self._dist_context = dist_context
        self._loss_sum = torch.zeros((), dtype=torch.float64, device="cuda")
        self._num_tokens = torch.zeros((), dtype=torch.float64, device="cuda")

    def build_forward_inputs(
        self, ctx: BuildForwardInputsContext[TextBatch]
    ) -> BuildForwardInputsResult[SequenceInput, SequenceHeadShared[SequenceCausalLMHeadShared], PerplexityState]:
        labels = ctx.batch.labels
        return BuildForwardInputsResult(
            input=SequenceInput(input_ids=ctx.batch.input_ids),
            shared=SequenceHeadShared(
                sequence=SequenceShared(position_ids=ctx.batch.position_ids),
                head=SequenceCausalLMHeadShared(labels=labels),
            ),
            state=PerplexityState(num_tokens=(labels != LM_IGNORE_INDEX).sum()),
        )

    def process_outputs(self, ctx: ProcessOutputsContext[SequenceCausalLMOutput, PerplexityState]):
        # The head returns the per-token cross-entropy loss. Padding positions hold zeros.
        self._loss_sum += ctx.pipeline_results.cross_entropy.sum()
        self._num_tokens += ctx.state.num_tokens

    def finalize(self, ctx: FinalizeContext) -> None:
        # Ranks that host no last pipeline stage add zeros, so a sum over all ranks covers every batch once.
        if self._dist_context.mesh_params.is_distributed:
            group = self._dist_context.mesh_for(FLAT_DOMAIN).get_group()
            dist.all_reduce(self._loss_sum, group=group)
            dist.all_reduce(self._num_tokens, group=group)

        mean_loss = (self._loss_sum / self._num_tokens).item()
        if self._dist_context.is_main_process:
            self._dist_context.logger.info(f"Mean loss {mean_loss:.4f}, perplexity {math.exp(mean_loss):.2f}")




def main():
    config = EvaluationConfig.model_validate(yaml.safe_load(Path(sys.argv[1]).read_text(encoding="utf-8")))
    # The checkpoint to evaluate, e.g. the original model or the export of a fine-tuning run.
    config.inference.model_stage_factory.source_checkpoint = Path(sys.argv[2])

    inference = InferenceConfigurator(
        mesh=config.mesh,
        parameters=config.inference,
        task_provider=lambda ctx: PerplexityTask(ctx.dist_context),
        model_provider=Qwen3Provider(config.model, lora=None),
        data_provider=AutoDataProvider(
            dataset_factory=lambda dist_context: build_dataset(config.data, dist_context),
            collator=TextDataset.collate,
            config=config.auto_data,
        ),
    ).configure()

    inference.infer()


if __name__ == "__main__":
    main()