Fine-tune Qwen3.8-27B on a single 8xB300 node
Preview
This feature is in Public Preview.
This example runs supervised fine-tuning (SFT) of Qwen/Qwen3.8-27B on UltraChat 200k for one full pass, on a single 8xB300 node. It uses plain PyTorch with Fully Sharded Data Parallel 2 (FSDP2) and no training framework.
Each B300 GPU has 288 GB of High Bandwidth Memory (HBM), enough for the example to pack conversations into 32,768-token blocks, four times the 8,192-token blocks an 8xH100 node fits for this model, so one 8xB300 node trains on as many tokens per step as four 8xH100 nodes.
The workload performs these steps:
- Uploads a snapshot of the local project with
code_source. - Builds the fused
causal-conv1dkernel for B300 and installs the libraries the model needs. - Launches one process per GPU with
torchrun. - Streams the dataset from the Hugging Face Hub and packs it so that no example attends to another.
Prerequisites
- The latest Databricks CLI installed and authenticated. See Install or update the Databricks CLI and Authentication for the Databricks CLI. Select and verify your workspace profile. See Get information about configuration profiles.
- Git installed and available on your
PATH. The CLI uses Git to package local code. - Access to 8xB300 for your workspace. 8xB300 is available under a waitlist. To request access, contact your Databricks account team.
The model and dataset are public, so no Hugging Face token is needed.
Project layout
Create a directory with these files:
Text
sft-qwen3.8-27b-b300/
├── workload.yaml # AI Runtime workload config (environment setup + torchrun launcher)
└── train.py # FSDP2 fine-tuning script
tip
To start from a runnable project instead of creating these files manually, clone the AI Runtime Cookbook and change to the recipe directory:
Bash
git clone https://github.com/databricks-solutions/ai-runtime-cookbook.git
cd ai-runtime-cookbook/language-models/sft-qwen3.8-27b-fsdp2-b300
Step 1: Write the workload YAML
workload.yaml requests one GPU_8xB300 node on the databricks_ai_v6 environment. The command builds causal-conv1d from source for
B300, installs pinned versions of transformers, flash-linear-attention, and
datasets, and launches train.py with 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 --watch
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
Step 2: Write the training script
train.py loads the pretrained checkpoint on every rank and shards each decoder layer with 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)
Examples are packed end to end into 32,768-token blocks. For each block, the script builds the signals that keep attention, the model's gated-delta recurrent state, and the loss from crossing an example boundary, and passes them to the model:
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,
)
The loss is computed on the assistant turns only. The complete script is listed in Full training script at the end of this page.
Step 3: Submit the run
Submit the workload and stream its logs:
Bash
databricks air run --file workload.yaml --watch
Step 4: Inspect the run
To check the run's status and view its logs, run:
Bash
databricks air get <job-run-id>
databricks air logs <job-run-id>
Rank 0 logs the training loss every 100 steps, then the total step count, time per step, and final loss.
On one 8xB300 node, a full pass takes 995 steps at about 8 seconds per step and ends with a training loss of about 0.74. The run takes about 2 hours 25 minutes end to end, including about 10 minutes of setup before training starts.
Full training script
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()