madokalif's picture
setup bundle: patch, lockfile, norm_stats, scripts
567a9bd verified
|
Raw
History Blame
3.24 kB

B200 migration bundle — pi0.5 LoRA SFT on turning_on_radio

Everything the B200 host needs that is not on the Hub. Full background is in ../HANDOFF.md ("MOVING B1 TRAINING TO ANOTHER HOST").

Contents

file what it is
openpi-sft.patch our diff vs upstream openpi (git apply after cloning)
openpi_commit.txt the base commit the patch applies to
pyproject.toml, uv.lock resolved dependency set, including the uv overrides
pip_freeze.txt exact versions of the environment that trained job 2136
check.py.fixed the transformers install gate, corrected to 4.53.2
assets/turning_on_radio/norm_stats.json do not recompute — 6.5h
train_b1_sft.sbatch, upload_ckpt_hf.py the launcher and the Hub uploader

Not in the bundle, copy separately:

  • data/b1k/turning_on_radio_v3.0/ — 4.9 GB
  • weights: pull a step-N from madokalif/pi05-b1k-turning-on-radio-lora (each includes optimizer.pt, so training resumes exactly, LR schedule included)

THE ONE THING THAT WILL BITE ON B200

The environment that produced these checkpoints runs torch 2.7.1+cu126, whose compiled arch list stops at sm_90:

['sm_50','sm_60','sm_70','sm_75','sm_80','sm_86','sm_90']

B200 is sm_100. This build has no kernel image for it and will fail at the first CUDA op. Install the cu128 build of the same torch version:

uv pip install torch==2.7.1 --index-url https://download.pytorch.org/whl/cu128

(torch-2.7.1+cu128-cp311-cp311-manylinux_2_28_x86_64.whl is published, so this is a build swap at an identical version — no API change.) Verify before launching:

python -c "import torch; print(torch.cuda.get_arch_list())"   # must contain sm_100

torchcodec==0.4.0 and torchvision are compiled against torch and may need matching rebuilds. torchcodec in particular was pinned to 0.4.0 only because 0.10.0 could not find this host's FFmpeg — re-test on B200 rather than assuming the pin carries over.

Steps

  1. git clone openpi at openpi_commit.txt, git apply openpi-sft.patch.
  2. Create the venv from pyproject.toml + uv.lock, then swap torch to cu128 (above) and uv pip install peft.
  3. Copy src/openpi/models_pytorch/transformers_replace/* into site-packages/transformers/, then overwrite transformers/models/siglip/check.py with check.py.fixed. Skipping this gives an AdaRMS state_dict mismatch — input_layernorm.dense.* missing, input_layernorm.weight unexpected — and the checkpoints will not load.
  4. Put norm_stats.json at outputs/assets/pi05_b1k_pytorch_local/turning_on_radio/.
  5. Edit dataset_root and pytorch_weight_path in the pi05_b1k_pytorch_local config (src/openpi/training/config.py) to the B200 paths.
  6. Launch train_b1_sft.sbatch, adjusting --gres / NGPU / BATCH. OPENPI_LORA=1 and OPENPI_LORA_RANK=32 are env vars, not flags.

Confirm it worked

Startup must log:

LoRA enabled (r=32): trainable 758.3M / 3.682B (20.60%)

A different fraction means the wrap landed somewhere else. On A6000 the run held 2.85 s/it at batch 16 across 4 GPUs with loss settling near 0.13 by step 6,000 — compare against that.