Faça o ajuste fino do Qwen3.8-27B em um nó 8xB300
Este notebook realiza a execução do ajuste fino supervisionado (SFT) do Qwen/Qwen3.8-27B no UltraChat 200k por uma passagem completa em um único nó 8xB300, com PyTorch simples e sua API Fully Sharded Data Parallel (FSDP2). Cada GPU B300 tem 288 GB de memória, então o exemplo empacota conversas em blocos de 32.768 tokens, quatro vezes o que um nó 8xH100 comporta para este modelo. Uma passagem completa tem 995 os passos e leva cerca de 2,5 horas, incluindo a configuração.
nota
Este exemplo requer o ambiente Databricks AI versão 6 (databricks_ai_v6) ou acima e acesso a aceleradores 8xB300.
info
Visualização
O treinamento em aceleradores 8xB300 (GPU_8xB300) está em Prévia Pública. A capacidade do B300 é limitada, portanto o acesso é restrito a clientes qualificáveis. Para solicitar acesso, entre em contato com sua equipe de account do Databricks.
Conectar ao compute de GPU serverless
- Clique no seletor de compute do notebook no canto superior direito e selecione GPU Serverless .
- No painel Ambiente à direita, selecione 8xB300 como Acelerador .
- Defina o ambiente base como AI v6 (
databricks_ai_v6). - Clique em Aplicar .
Instalar dependências
Instale as bibliotecas que o modelo precisa. O kernel causal-conv1d é compilado a partir do código-fonte para B300.
Python
import os
os.environ["TORCH_CUDA_ARCH_LIST"] = "10.0"
os.environ["CAUSAL_CONV1D_FORCE_BUILD"] = "TRUE"
Python
%pip install --no-cache-dir ninja==1.13.0 packaging==25.0
Python
%pip install --no-cache-dir --no-build-isolation git+https://github.com/Dao-AILab/causal-conv1d.git@v1.6.2
Python
%pip install --no-cache-dir transformers==5.16.1 flash-linear-attention==0.5.2 datasets==4.8.5
Python
dbutils.library.restartPython()
Configurar modelo e dataset
As revisões de modelos e datasets são pin para que as execuções permaneçam comparáveis conforme os repository Hugging Face se movem. Ambos são públicos, portanto não é necessário um access token.
Python
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
Definir o dataset SFT empacotado
Cada GPU faz transmissão de seu próprio fragmento do UltraChat e empacota conversas em blocos de 32.768 tokens. A perda é compute apenas nos turnos do assistente, e nenhuma conversa atenção a outra dentro de um bloco.
Python
import os
import torch
from torch.utils.data import IterableDataset, get_worker_info
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):
from datasets import load_dataset
from datasets.distributed import split_dataset_by_node
# 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 = int(os.environ["RANK"]) * num_workers + worker_id
shard_world = int(os.environ["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)
Definir o modelo e o loop de treinamento
Cada GPU carrega o checkpoint pré-treinado e o FSDP2 o particiona entre as 8 GPUs. A classificação 0 registra a perda de treinamento a cada 100 os passos.
Python
import logging
import time
def build_sharded_model(log) -> 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.)
"""
from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
from transformers import AutoModelForImageTextToText
# 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 int(os.environ["RANK"]) == 0:
log.info(
f"model built + weights loaded: {model.config.model_type}, {len(layers)} layers, "
f"sharded over {os.environ['WORLD_SIZE']}"
)
return model
def run_training_loop(model: torch.nn.Module, data, log) -> None:
import torch.distributed as dist
rank = int(os.environ["RANK"])
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}"
)
Iniciar nas 8 GPUs do nó
O decorador @distributed de serverless_gpu executa a função uma vez por GPU no iniciador distribuído do PyTorch, portanto, cada processo obtém seu próprio RANK, LOCAL_RANK e WORLD_SIZE. timeout=None desativa o tempo limite local default de 3 horas do decorador, o que deixa pouca margem para uma execução de cerca de 2,5 horas.
Python
from serverless_gpu import distributed
@distributed(gpus=8, gpu_type="B300", timeout=None)
def run_sft():
import torch.distributed as dist
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
local_rank = int(os.environ["LOCAL_RANK"])
logging.Formatter.converter = time.gmtime # UTC timestamps
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",
force=True, # the notebook's launcher already configures logging at WARNING
)
log = logging.getLogger("sft")
torch.cuda.set_device(local_rank)
torch.manual_seed(0)
dist.init_process_group("nccl")
try:
model = build_sharded_model(log)
tokenizer = AutoTokenizer.from_pretrained(MODEL, revision=MODEL_REVISION)
if int(os.environ["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, log)
finally:
if dist.is_initialized():
dist.destroy_process_group()
run_sft.distributed()