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.
- 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
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
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()