Script de treinamento completo

info

Visualização

Esse recurso está em Prévia Pública.

Este exemplo executa uma execução de ajuste fino supervisionado (SFT) do Qwen/Qwen3.8-27B no UltraChat 200k por uma passagem completa, em um nó 8xB300 único. Ele usa PyTorch simples com Fully Sharded Data Parallel 2 (FSDP2) e nenhum framework de treinamento.

Cada GPU B300 possui 288 GB de HBM (High Bandwidth Memory), o suficiente para o exemplo empacotar conversas em blocos de 32.768 tokens, quatro vezes os blocos de 8.192 tokens que um nó 8xH100 cabe para este modelo, de modo que um nó 8xB300 ensinar em tantos tokens por o passo quanto quatro nós 8xH100.

A carga de trabalho executa os passos:

• Faz upload de um snapshot do projeto local com code_source.
• Cria o kernel causal-conv1d fundido para B300 e instala as bibliotecas de que o modelo precisa.
• Inicia um processo por GPU com torchrun.
• Transmite o dataset do Hugging Face Hub e o empacota para que nenhum exemplo participe de outro.

Pré-requisitos

• A versão mais recente da CLI do Databricks instalada e autenticada. Consulte Instalar ou atualizar a CLI do Databricks e Autenticação para a CLI do Databricks.
• Git instalado e disponível em seu PATH. A CLI usa Git para criar um pacote com o código local.
• Acesso a 8xB300 para o seu workspace. 8xB300 está disponível mediante lista de espera. Para solicitar acesso, entre em contato com sua equipe de account do Databricks.

O modelo e o dataset são públicos, portanto, não é necessário um token do Hugging Face.

Disposição do projeto

Crie um diretório com estes arquivos:

Text

sft-qwen3.8-27b-b300/
├── workload.yaml       # AI Runtime workload config (environment setup + torchrun launcher)
└── train.py            # FSDP2 fine-tuning script

O passo 1: Escrever o YAML da carga de trabalho

workload.yaml solicita um nó GPU_8xB300 no ambiente databricks_ai_v6. O comando compila causal-conv1d a partir da fonte para B300, instala versões pin de transformers, flash-linear-attention e datasets e inicia train.py com torchrun:

YAML

# Qwen3.8-27B SFT fine-tuning on one 8xB300 node -- plain torch, no training framework.
#
# Self-contained torch + HF transformers + FSDP2 training loop (launched via
# torchrun): load pretrained weights -> shard + data prep + one full pass over the
# dataset + teardown. The training loop reports its wall time and time per step on rank 0.
#
# Model:  HF transformers Qwen3.8-27B (VLM-style model_type qwen3_5). Every rank loads
#         the pretrained checkpoint and FSDP2 shards it across ranks (see train.py). Model
#         and dataset revisions are pinned so runs stay comparable as HF repos move.
# Data:   UltraChat 200k (SFT) streamed from the HF hub, chat-templated + tokenized
#         with the Qwen3.8 tokenizer into seq-32768 blocks (response-only loss).
# Image:  databricks_ai_v6 -- the "v6" AI Runtime environment version (a CUDA 13 base) with a
#         Blackwell-ready torch and FlashAttention-2. We build causal-conv1d from source against
#         that torch, and install transformers + fla.
#
# Run it directly on a workspace with 8xB300 access:
#   databricks air run -f workload.yaml

experiment_name: air-b300-qwen3-8-27b-sft

environment:
  version: databricks_ai_v6
  dependencies: []

compute:
  accelerator_type: GPU_8xB300
  num_accelerators: 8

env_variables:
  PYTORCH_ALLOC_CONF: "expandable_segments:True"

code_source:
  # Snapshot this dir so the launched job has train.py.
  root_path: "."

max_retries: 0
timeout_minutes: 240

command: |
  set -eux
  cd "$CODE_SOURCE_PATH"

  export HF_HOME=/tmp/.hf
  export NCCL_DEBUG=WARN
  # Build the kernels below for Blackwell (sm_100) only.
  export TORCH_CUDA_ARCH_LIST="10.0"

  # === ---- Env Set up ---------------------------------------------------- ===
  #
  echo "Env Setup"

  # The fused causal-conv1d (for the gated-delta layers), built from the upstream Dao-AILab
  # source against the environment's torch. torch and FlashAttention-2 come from the image.
  python -m pip install --no-cache-dir ninja==1.13.0 packaging==25.0

  # causal_conv1d, tag v1.6.2.
  git clone --depth 1 --branch v1.6.2 https://github.com/Dao-AILab/causal-conv1d.git /tmp/causal-conv1d
  CAUSAL_CONV1D_FORCE_BUILD=TRUE python -m pip install --no-cache-dir --no-build-isolation /tmp/causal-conv1d

  # transformers with Qwen3.8 support (newer than the image's baked version).
  python -m pip install --no-cache-dir transformers==5.16.1
  # flash-linear-attention: the fused gated-delta kernel the model uses (installing it is enough).
  python -m pip install --no-cache-dir flash-linear-attention==0.5.2
  # datasets: train.py streams the training data from the HF hub.
  python -m pip install --no-cache-dir datasets==4.8.5

  # === ---- Run Main Training Loop ---------------------------------------- ===
  #
  echo "Run Main Training Loop"

  # One process per GPU on the single node.
  python -m torch.distributed.run --nproc_per_node=8 train.py

o passo 2: Escreva o script de treinamento

train.py carrega o checkpoint pré-treinado em cada rank e fragmenta cada camada de decodificador com FSDP2:

Python

model = AutoModelForImageTextToText.from_pretrained(
    MODEL,
    revision=MODEL_REVISION,
    attn_implementation="flash_attention_2",
    dtype=torch.bfloat16,
)
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})

mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
for layer in model.model.language_model.layers:
    fully_shard(layer, mp_policy=mp)
fully_shard(model, mp_policy=mp)

Os exemplos são empacotados de ponta a ponta em blocos de 32.768 tokens. Para cada bloco, o script cria os sinais que mantêm a atenção, o estado recorrente delta-gated do modelo e a perda de cruzar um limite de exemplo, e os passa para o modelo:

Python

out = model(
    input_ids=ids,
    labels=labels,
    position_ids=position_ids,
    cu_seq_lens_q=cu_seqlens,
    cu_seq_lens_k=cu_seqlens,
    max_length_q=max_len,
    max_length_k=max_len,
    seq_idx=seq_idx,
)

A perda é computada apenas nas respostas do assistente. Passo 2: Escreva o script de treinamento

O passo 3: Enviar a execução

Envie a carga de trabalho e transmissão seus Logs:

Bash

databricks air run --file workload.yaml --watch

O passo 4: Inspecionar a execução

Para verificar o status da execução e view seus logs, execute:

Bash

databricks air get <job-run-id>
databricks air logs <job-run-id>

O rank 0 Logs a perda de treinamento a cada 100 passos, depois o número total de passos, o tempo por passo e a perda final.

Em um nó 8xB300, uma execução completa leva 995 passos a cerca de 8 segundos por passo e termina com uma perda de treinamento de cerca de 0,74. A execução leva cerca de 2 horas e 25 minutos de ponta a ponta, incluindo cerca de 10 minutos de configuração antes de começar o treinamento.

Script de treinamento completo

Python

#!/usr/bin/env python3
"""Qwen3.8-27B SFT fine-tuning example -- self-contained, plain torch (no training framework).

A torch + HF-transformers + FSDP2 training loop, one process per GPU via torchrun. The phases:

  - build:  load the pretrained Qwen3.8-27B checkpoint on every rank, then FSDP2-shard it
  - data:   stream UltraChat 200k (SFT), chat-template + tokenize with a response-only loss
            mask (loss on the assistant turns only), and pack into fixed seq-len blocks
  - train:  one full pass over the dataset

Packing is block-diagonal: examples are concatenated into SEQ_LEN blocks, but no token
attends across an example boundary. _pack_block derives the per-block signals (position_ids,
cu_seqlens, seq_idx) that keep attention, the gated-delta recurrent state, and the loss from
crossing boundaries -- which needs a flash-attention varlen backend (see workload.yaml).

See workload.yaml for the matching environment.
"""

import logging
import os
import time

import torch
import torch.distributed as dist
from torch.utils.data import DataLoader, IterableDataset, get_worker_info

from datasets import load_dataset
from datasets.distributed import split_dataset_by_node
from transformers import AutoModelForImageTextToText, AutoTokenizer

from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard

RANK = int(os.environ["RANK"])
LOCAL_RANK = int(os.environ["LOCAL_RANK"])
WORLD_SIZE = int(os.environ["WORLD_SIZE"])

# Use UTC instead of local time
logging.Formatter.converter = time.gmtime

logging.basicConfig(
    level=logging.INFO,
    format=f"[RANK: {LOCAL_RANK}] %(asctime)s %(levelname)s %(name)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)

log = logging.getLogger(__name__)


MODEL = "Qwen/Qwen3.8-27B"
MODEL_REVISION = "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0"
# SFT dataset streamed from the HF hub (UltraChat 200k, MIT-licensed multi-turn chat).
# Revision pinned so runs stay comparable as the repo moves.
DATASET = "HuggingFaceH4/ultrachat_200k"
DATASET_REVISION = "8049631c405ae6576f93f445c6b8166f76f5505a"
SPLIT = "train_sft"
# Tokens per packed block. B300's 288 GB of HBM fits 4x the H100 recipe's 8192, so one
# 8xB300 node takes as many tokens per step as four 8xH100 nodes.
SEQ_LEN = 32768


class PackedTokenDataset(IterableDataset):
    """Streams this rank's shard of the dataset and packs it into [1, SEQ_LEN] blocks.

    Each rank streams a disjoint shard (split_dataset_by_node), SFT-tokenizes examples
    (chat template + response-only loss mask), and jams them end-to-end into a buffer;
    whenever the buffer overfills SEQ_LEN it chops off one block and _pack_block builds
    the block-diagonal signals + masked labels. Nothing is materialized up front.

    Each item is (block, ran_out): ran_out is 0 for a real block, then 1 on a final all-masked
    block once the shard is exhausted.
    """

    def __init__(self, tokenizer, device: torch.device):
        self.tok = tokenizer
        self.device = device

    def _tokenize_sft(self, messages: list[dict]) -> tuple[list[int], list[bool]]:
        """Chat-template a (possibly multi-turn) conversation and return its token ids plus a
        per-token mask that is True only on the assistant turns' response tokens (SFT loss).

        Each assistant turn is located by prefix length: templating up to just before the turn
        (add_generation_prompt=True) vs through it brackets that turn's response tokens. This
        works with any chat template, without relying on {% generation %} markers.
        """
        full_ids = self.tok.apply_chat_template(messages, tokenize=True, return_dict=False)
        loss_mask = [False] * len(full_ids)
        for i, msg in enumerate(messages):
            if msg["role"] != "assistant":
                continue
            prefix = self.tok.apply_chat_template(
                messages[:i], add_generation_prompt=True, tokenize=True, return_dict=False
            )
            through = self.tok.apply_chat_template(messages[: i + 1], tokenize=True, return_dict=False)
            for j in range(len(prefix), min(len(through), len(full_ids))):
                loss_mask[j] = True  # this assistant turn's response tokens contribute to loss
        return full_ids, loss_mask

    def _pack_block(self, ids: list[int], interior_starts: list[int], loss_mask: list[bool]):
        """Turn one SEQ_LEN block of token ids + the offsets where interior examples begin +
        the per-token SFT loss mask into (input_ids, labels, position_ids, cu_seqlens, seq_idx, max_len).

        Block-diagonal, so nothing mixes across an example boundary:
          * position_ids -- reset to 0 per segment
          * cu_seqlens   -- segment boundaries, for the full-attention layers
          * seq_idx      -- per-token segment id, for the gated-delta layers
          * labels       -- -100 on prompt tokens (SFT loss is response-only) and on each
                            segment's first token (so the causal-shifted loss can't cross a boundary)
        """
        bounds = [0, *sorted({s for s in interior_starts if 0 < s < SEQ_LEN}), SEQ_LEN]
        ids_t = torch.tensor(ids, dtype=torch.long)
        position_ids = torch.empty(SEQ_LEN, dtype=torch.long)
        seq_idx = torch.empty(SEQ_LEN, dtype=torch.int32)
        labels = ids_t.clone()
        labels[~torch.tensor(loss_mask, dtype=torch.bool)] = -100  # SFT: mask prompt tokens
        for i, (a, b) in enumerate(zip(bounds[:-1], bounds[1:])):
            position_ids[a:b] = torch.arange(b - a)
            seq_idx[a:b] = i
            if a > 0:
                labels[a] = -100  # causal-shifted target would come from the previous segment
        cu_seqlens = torch.tensor(bounds, dtype=torch.int32)
        max_len = max(b - a for a, b in zip(bounds[:-1], bounds[1:]))
        return (
            ids_t.unsqueeze(0).to(self.device),
            labels.unsqueeze(0).to(self.device),
            position_ids.unsqueeze(0).to(self.device),
            cu_seqlens.to(self.device),
            seq_idx.unsqueeze(0).to(self.device),
            max_len,
        )

    def _shard_examples(self):
        # Shard across every (rank, DataLoader-worker) stream, not just per rank, so
        # num_workers>1 gives each worker a disjoint slice instead of replaying the whole
        # rank shard. With the default num_workers=1 this reduces to per-rank sharding.
        worker = get_worker_info()
        worker_id = worker.id if worker is not None else 0
        num_workers = worker.num_workers if worker is not None else 1
        shard_rank = RANK * num_workers + worker_id
        shard_world = WORLD_SIZE * num_workers
        ds = load_dataset(DATASET, split=SPLIT, streaming=True, revision=DATASET_REVISION)
        ds = split_dataset_by_node(ds, rank=shard_rank, world_size=shard_world)
        yield from ds

    def __iter__(self):
        buf: list[int] = []  # token ids, examples jammed end-to-end
        loss: list[bool] = []  # parallel SFT mask: True where the token contributes to loss
        starts: list[int] = []  # offsets in buf where each example begins
        for ex in self._shard_examples():
            starts.append(len(buf))
            ids, ex_loss = self._tokenize_sft(ex["messages"])
            buf.extend(ids)
            loss.extend(ex_loss)
            while len(buf) >= SEQ_LEN:
                yield self._pack_block(buf[:SEQ_LEN], starts, loss[:SEQ_LEN]), torch.tensor(0.0)
                buf = buf[SEQ_LEN:]
                loss = loss[SEQ_LEN:]
                # keep only example starts landing in the carried-over remainder
                starts = [s - SEQ_LEN for s in starts if s - SEQ_LEN > 0]
        yield self._pack_block([0] * SEQ_LEN, [], [False] * SEQ_LEN), torch.tensor(1.0)


def build_sharded_model() -> torch.nn.Module:
    """Load the pretrained Qwen3.8-27B on every rank, then FSDP2-shard each decoder layer + root.

    Every rank loads the full checkpoint, fully_shard reshards it in place, and .cuda() moves
    each rank's shard to GPU. Simple, at the cost of every rank holding the whole model in host
    RAM transiently -- fine at 27B. (For a much larger model, load on rank 0 and broadcast each
    rank its shard instead.)
    """
    # Explicit flash_attention_2: the packing needs a varlen flash-attention backend, and
    # transformers otherwise defaults to sdpa. FlashAttention-3 is Hopper-only, so B300 uses
    # the FlashAttention-2 build that ships in the databricks_ai_v6 environment.
    model = AutoModelForImageTextToText.from_pretrained(
        MODEL,
        revision=MODEL_REVISION,
        attn_implementation="flash_attention_2",
        dtype=torch.bfloat16,
    )
    # The KV cache only helps generation; in training it wastes memory and is incompatible
    # with gradient checkpointing. (text_config is this VLM-style model's language tower.)
    model.config.use_cache = False
    model.config.text_config.use_cache = False
    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})

    mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
    layers = model.model.language_model.layers
    for layer in layers:
        fully_shard(layer, mp_policy=mp)
    fully_shard(model, mp_policy=mp)
    model.cuda()

    if RANK == 0:
        log.info(
            f"model built + weights loaded: {model.config.model_type}, {len(layers)} layers, sharded over {WORLD_SIZE}"
        )
    return model


def run_training_loop(model: torch.nn.Module, data) -> None:
    opt = torch.optim.AdamW(model.parameters(), lr=1e-5, betas=(0.9, 0.95), weight_decay=0.0, fused=True)
    model.train()

    loss = None
    loop_start = time.perf_counter()
    for step, (batch, ran_out) in enumerate(data):
        exhausted = bool(ran_out)  # this rank's own flag; its final all-masked block has a NaN loss
        # Ranks get different numbers of batches, so they must stop on the same step: a rank that
        # stopped early would hang the others in the next collective.
        ran_out = ran_out.cuda(non_blocking=True)
        dist.all_reduce(ran_out, op=dist.ReduceOp.MAX)

        ids, labels, position_ids, cu_seqlens, seq_idx, max_len = batch
        # The loader builds batches on CPU (a worker can't return CUDA tensors); move to GPU here.
        ids = ids.cuda(non_blocking=True)
        labels = labels.cuda(non_blocking=True)
        position_ids = position_ids.cuda(non_blocking=True)
        cu_seqlens = cu_seqlens.cuda(non_blocking=True)
        seq_idx = seq_idx.cuda(non_blocking=True)
        out = model(
            input_ids=ids,
            labels=labels,
            position_ids=position_ids,
            cu_seq_lens_q=cu_seqlens,
            cu_seq_lens_k=cu_seqlens,
            max_length_q=max_len,
            max_length_k=max_len,
            seq_idx=seq_idx,
        )
        out.loss.backward()
        opt.step()
        opt.zero_grad(set_to_none=True)
        if not exhausted:
            loss = out.loss
        if (step + 1) % 100 == 0 and RANK == 0:
            log.info(f"[train] step {step + 1}, loss={loss.item():.3f}")

        # .item() is a sync point, so check it at the end rather than before the step.
        if ran_out.item():
            break
    torch.cuda.synchronize()  # let all queued GPU work finish before we stop the clock
    elapsed = time.perf_counter() - loop_start
    if RANK == 0:
        log.info(
            f"[train] {step + 1} steps in {elapsed:.2f}s ({elapsed / (step + 1):.2f}s/step), final loss={loss.item():.3f}"
        )


def main() -> None:
    torch.cuda.set_device(LOCAL_RANK)
    torch.manual_seed(0)
    dist.init_process_group("nccl")
    try:
        model = build_sharded_model()
        tokenizer = AutoTokenizer.from_pretrained(MODEL, revision=MODEL_REVISION)
        if RANK == 0:
            log.info(f"streaming dataset={DATASET}, seq_len={SEQ_LEN}")

        data = DataLoader(
            PackedTokenDataset(tokenizer, torch.device("cpu")),
            batch_size=None,
            num_workers=1,
            prefetch_factor=4,
            pin_memory=True,
        )

        run_training_loop(model, data)
    finally:
        if dist.is_initialized():
            dist.destroy_process_group()


if __name__ == "__main__":
    main()

O script completo está listado em Full training script no final desta página. recursos

• Usar a CLI do Databricks com AI Runtime
• Referência de YAML de Workload
• Cookbook do AI Runtime