Fine-tune Qwen3.8-27B on one 8xB300 node
This notebook runs supervised fine-tuning (SFT) of Qwen/Qwen3.8-27B on UltraChat 200k for one full pass on a single 8xB300 node, with plain PyTorch and its Fully Sharded Data Parallel API (FSDP2). Each B300 GPU has 288 GB of memory, so the example packs conversations into 32,768-token blocks, four times what an 8xH100 node fits for this model. A full pass is 995 steps and takes about 2.5 hours, including setup.
note
This example requires the Databricks AI environment version 6 (databricks_ai_v6) or above and access to 8xB300 accelerators.
Preview
Training on 8xB300 (GPU_8xB300) accelerators is in Public Preview. B300 capacity is limited, so access is gated to qualifying customers. To request access, contact your Databricks account team.
Connect to serverless GPU compute
- Click the notebook's compute selector in the upper right and select Serverless GPU.
- In the Environment panel on the right, select 8xB300 as the Accelerator.
- Set the base environment to AI v6 (
databricks_ai_v6). - Click Apply.
Install dependencies
Install the libraries the model needs. The causal-conv1d kernel is built from source for 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()
Configure model and dataset
Model and dataset revisions are pinned so runs stay comparable as the Hugging Face repositories move. Both are public, so no access token is needed.
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
Define the packed SFT dataset
Each GPU streams its own shard of UltraChat and packs conversations into 32,768-token blocks. The loss is computed on the assistant turns only, and no conversation attends to another within a block.
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)
Define the model and training loop
Every GPU loads the pretrained checkpoint, and FSDP2 shards it across the 8 GPUs. Rank 0 logs the training loss every 100 steps.
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}"
)
Launch on the node's 8 GPUs
The @distributed decorator from serverless_gpu runs the function one time per GPU under PyTorch's distributed launcher, so each process gets its own RANK, LOCAL_RANK, and WORLD_SIZE. timeout=None turns off the decorator's default 3-hour local timeout, which leaves little headroom for a run of about 2.5 hours.
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()