Ajustez Qwen3.8-27B sur un nœud 8xB300
Ce notebook exécute l'affinement supervisé (SFT) de Qwen/Qwen3.8-27B sur UltraChat 200k pour un passage complet sur un seul nœud 8xB300, avec du PyTorch brut et son API Fully Sharded Data Parallel (FSDP2). Chaque GPU B300 dispose de 288 Go de mémoire, de sorte que l'exemple compresse les conversations en blocs de 32 768 jetons, soit quatre fois plus qu'un nœud 8xH100 pour ce modèle. Un passage complet représente 995 étapes et prend environ 2,5 heures, configuration comprise.
remarque
Cet exemple nécessite l’environnement Databricks AI version 6 (databricks_ai_v6) ou supérieure et l’accès aux accélérateurs 8xB300.
info
Aperçu
L'entraînement sur les accélérateurs 8xB300 (GPU_8xB300) est en version préliminaire publique. La capacité B300 est limitée, l'accès est donc réservé aux clients éligibles. Pour demander l’accès, contactez l’équipe chargée de votre compte Databricks.
Connexion au compute GPU serverless
- Cliquez sur le sélecteur de compute du notebook en haut à droite et sélectionnez GPU serverless .
- Dans le panneau Environnement à droite, sélectionnez 8xB300 comme Accélérateur .
- Définir l’environnement de base sur AI v6 (
databricks_ai_v6). - Cliquez sur Appliquer .
Installer les dépendances
Installez les bibliothèques dont le modèle a besoin. Le noyau causal-conv1d est compilé à partir de la source pour 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()
Configurer le modèle et le dataset
Les révisions des modèles et des datasets sont pin afin que les exécutions restent comparables à mesure que les repository Hugging Face changent. Les deux sont publics, aucun jeton d'accès n'est donc nécessaire.
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
Définir le dataset SFT compressé
Chaque GPU Stream sa propre partition d’UltraChat et regroupe les conversations en blocs de 32 768 jetons. La perte est calculée uniquement sur les tours de l’assistant, et aucune conversation ne prête attention à une autre au sein d’un bloc.
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)
Définir le modèle et la boucle d'entraînement
Chaque GPU charge le point de contrôle pré-entraîné, et FSDP2 le fragmente sur les 8 GPU. Le rang 0 Logs la perte d’entraînement toutes les 100 étapes.
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}"
)
Lancez sur les 8 GPU du nœud
Le décorateur @distributed de serverless_gpu exécute la fonction une fois par GPU sous le lanceur distribué de PyTorch, de sorte que chaque processus obtient ses propres RANK, LOCAL_RANK et WORLD_SIZE. timeout=None désactive le délai d'expiration local default de 3 heures du décorateur, ce qui laisse peu de marge pour une exécution d'environ 2,5 heures.
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()