Massed Compute

Launch GPU VM

# Launch via Massed Compute dashboard or API
# Product: gpu_2x_h100
# Image: 184 (Ubuntu Server 24.04 w/ Drivers)
# SSH Key: attach your public key
# Instance name: deepspeed-multigpu

Wait until the VM is running and copy the SSH details.

Verify GPU Access

ssh -o IdentitiesOnly=yes Ubuntu@YOUR_VM_IP \
  'nvidia-smi -L; nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader; nvidia-smi topo -m'

This run printed two lines: NVIDIA H100 PCIe, 81559 MiB, 580.126.16. nvidia-smi -L must show 2 GPUs. topo -m printed PHB between GPU0 and GPU1 — PCIe through the CPU host bridge, not NVLink (NV#). Training still used both cards.

Lock the firewall to SSH only. No public training UI is required.

ssh -o IdentitiesOnly=yes Ubuntu@YOUR_VM_IP 'bash -s' <<'EOF'
sudo ufw default deny incoming
sudo ufw default allow outgoing
sudo ufw allow 22/tcp
echo y | sudo ufw enable
sudo ufw status
EOF

Install DeepSpeed

Pin whatever uv pip actually installs. This VM printed DeepSpeed 0.19.6, torch 2.12.1+cu130, Transformers 5.16.1, Accelerate 1.14.0. DS_BUILD_OPS=0 skips JIT CUDA ops (FusedAdam, transformer kernels). ZeRO does not need those ops. import deepspeed then warns CUDA_HOME does not exist — expected on image 184 without a CUDA toolkit.

ssh -o IdentitiesOnly=yes Ubuntu@YOUR_VM_IP 'bash -s' <<'EOF'
set -euxo pipefail
sudo apt-get update -qq
sudo DEBIAN_FRONTEND=noninteractive apt-get install -y python3.12-venv curl ninja-build
curl -LsSf https://astral.sh/uv/install.sh | sh
export PATH="$HOME/.local/bin:$PATH"
export UV_TORCH_BACKEND=cu130
uv venv --python 3.12 ~/ds_env
source ~/ds_env/bin/activate
uv pip install -U pip packaging setuptools wheel ninja py-cpuinfo
uv pip install torch==2.12.1 torchvision
DS_BUILD_OPS=0 uv pip install transformers datasets accelerate deepspeed
python -c "import torch,deepspeed,transformers; print(torch.__version__, deepspeed.__version__, transformers.__version__, torch.cuda.device_count())"
EOF

Stop if device_count is not 2.

Train script and ZeRO JSON

Copy train_ds.py, ds_z2.json, and ds_z3.json to the VM. The script accepts --local_rank because the DeepSpeed launcher injects it.

train_ds.py (HF Trainer, causal LM, optional --deepspeed):

#!/usr/bin/env python3
from __future__ import annotations
import argparse
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM, AutoTokenizer, DataCollatorForLanguageModeling,
    Trainer, TrainingArguments,
)

def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="Qwen/Qwen2.5-0.5B-Instruct")
    p.add_argument("--output_dir", required=True)
    p.add_argument("--deepspeed")
    p.add_argument("--max_steps", type=int, default=20)
    p.add_argument("--seq_len", type=int, default=256)
    p.add_argument("--per_device_train_batch_size", type=int, default=1)
    p.add_argument("--gradient_accumulation_steps", type=int, default=4)
    p.add_argument("--rows", type=int, default=1024)
    p.add_argument("--lr", type=float, default=2e-5)
    p.add_argument("--gradient_checkpointing", action="store_true")
    p.add_argument("--local_rank", type=int, default=-1)
    args, _unknown = p.parse_known_args()

    tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    raw = load_dataset("yahma/alpaca-cleaned", split=f"train[:{args.rows}]")

    def to_text(ex):
        instr, inp, out = ex["instruction"], ex.get("input") or "", ex["output"]
        if inp:
            text = f"### Instruction:\n{instr}\n\n### Input:\n{inp}\n\n### Response:\n{out}"
        else:
            text = f"### Instruction:\n{instr}\n\n### Response:\n{out}"
        return {"text": text}

    ds = raw.map(to_text, remove_columns=raw.column_names)
    ds = ds.map(lambda b: tok(b["text"], truncation=True, max_length=args.seq_len, padding=False),
                batched=True, remove_columns=["text"])
    model = AutoModelForCausalLM.from_pretrained(
        args.model, dtype=torch.bfloat16, trust_remote_code=True, attn_implementation="sdpa",
    )
    if args.gradient_checkpointing:
        model.gradient_checkpointing_enable()
        model.config.use_cache = False
    targs = TrainingArguments(
        output_dir=args.output_dir, max_steps=args.max_steps,
        per_device_train_batch_size=args.per_device_train_batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.lr, bf16=True, logging_steps=1,
        save_strategy="steps", save_steps=args.max_steps, save_total_limit=1,
        report_to=[], dataloader_num_workers=0, deepspeed=args.deepspeed,
        gradient_checkpointing=args.gradient_checkpointing,
        remove_unused_columns=False, ddp_find_unused_parameters=False,
        logging_nan_inf_filter=False,
    )
    trainer = Trainer(model=model, args=targs, train_dataset=ds,
                      data_collator=DataCollatorForLanguageModeling(tok, mlm=False),
                      processing_class=tok)
    trainer.train()
    trainer.save_model(args.output_dir)
    tok.save_pretrained(args.output_dir)

if __name__ == "__main__":
    main()

ds_z2.json:

{
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "bf16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "reduce_scatter": true,
    "overlap_comm": false,
    "contiguous_gradients": true
  },
  "gradient_clipping": 1.0,
  "steps_per_print": 1,
  "wall_clock_breakdown": false
}

ds_z3.json — same auto batch keys, "stage": 3, plus "stage3_gather_16bit_weights_on_model_save": true so save_model writes a single model.safetensors.

Do not add "optimizer": { "type": "AdamW" } unless you have CUDA_HOME or set "torch_adam": true. On this image, that JSON made DeepSpeed JIT FusedAdam and exited with CUDA_HOME environment variable is not set. Let HF Trainer own Adam.

0.5B one-GPU baseline

ssh -o IdentitiesOnly=yes Ubuntu@YOUR_VM_IP 'bash -s' <<'EOF'
source ~/ds_env/bin/activate
CUDA_VISIBLE_DEVICES=0 python -u ~/train_ds.py \
  --output_dir ~/ds-out/0p5-1gpu --max_steps 20 --seq_len 256
EOF

This run: train_runtime 15.76 s, loss 1.368, peak 5985 MiB on GPU 0, 0 MiB on GPU 1. Wrote ~/ds-out/0p5-1gpu/model.safetensors (943 MB / 494M params). Full fine-tune, not a LoRA adapter.

0.5B two-GPU DDP (no DeepSpeed)

ssh -o IdentitiesOnly=yes Ubuntu@YOUR_VM_IP \
  'source ~/ds_env/bin/activate && torchrun --nproc_per_node=2 --standalone ~/train_ds.py --output_dir ~/ds-out/0p5-ddp --max_steps 20 --seq_len 256'

This run: train_runtime 18.08 s, 8.848 samples/s (about 1.7× the 1-GPU 5.077 samples/s), peak 7257 / 7033 MiB, 100% util on both. GPU 1 is no longer idle. Same peak as the August 31 capture.

0.5B ZeRO-2 and ZeRO-3

The DeepSpeed launcher warns Unable to find hostfile, will proceed with training with local resources only on a single node. That is expected. Do not invent a hostfile for one VM.

ssh -o IdentitiesOnly=yes Ubuntu@YOUR_VM_IP 'bash -s' <<'EOF'
source ~/ds_env/bin/activate
deepspeed --num_gpus=2 ~/train_ds.py --deepspeed ~/ds_z2.json \
  --output_dir ~/ds-out/0p5-z2 --max_steps 20 --seq_len 256
deepspeed --num_gpus=2 ~/train_ds.py --deepspeed ~/ds_z3.json \
  --output_dir ~/ds-out/0p5-z3 --max_steps 20 --seq_len 256
EOF

On 0.5B, ZeRO uses more VRAM and fewer samples/s than DDP. The overhead is real. Size the SKU for the 7B job.

7B full FT — why ZeRO exists

Same script, --gradient_checkpointing, 8 steps, Qwen/Qwen2.5-7B-Instruct.

One GPU (CUDA_VISIBLE_DEVICES=0, default GAS 4):

ssh -o IdentitiesOnly=yes Ubuntu@YOUR_VM_IP 'bash -s' <<'EOF'
source ~/ds_env/bin/activate
CUDA_VISIBLE_DEVICES=0 python -u ~/train_ds.py \
  --model Qwen/Qwen2.5-7B-Instruct \
  --output_dir ~/ds-out/7b-1gpu \
  --max_steps 8 --seq_len 256 --gradient_checkpointing
EOF

This run fit: train_runtime 55.43 s, loss 1.179, peak 61601 MiB on GPU 0, GPU 1 idle. Gradient checkpointing is why 7B full FT landed under 80 GB on one H100. train_runtime includes the 15 GB shard write, so wall time moves with disk; VRAM matched the August 31 capture exactly.

ZeRO-2 on both cards OOM’d. First step tried to allocate another 14.19 GiB (torch.zeros_like for Adam exp_avg_sq) with ~65 GB already in use on each GPU. ZeRO-2 still replicates parameters. Do not treat ZeRO-2 as a free 2× memory cut on this HF Trainer + torch Adam path.

ZeRO-3 completed:

ssh -o IdentitiesOnly=yes Ubuntu@YOUR_VM_IP 'bash -s' <<'EOF'
source ~/ds_env/bin/activate
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
deepspeed --num_gpus=2 ~/train_ds.py \
  --model Qwen/Qwen2.5-7B-Instruct \
  --deepspeed ~/ds_z3.json \
  --output_dir ~/ds-out/7b-z3 \
  --max_steps 8 --seq_len 256 --gradient_checkpointing \
  --gradient_accumulation_steps 1
EOF

This run: train_runtime 125.2 s, then a second independent launch 134.8 s, loss 1.157 both times, peak 68855 / 68851 MiB both times, 100% util on both. Wrote ~/ds-out/7b-z3/model.safetensors (15 GB / 7.616B params) and checkpoint-8/zero_to_fp32.py. nvidia-smi sat at ~68 GB for the whole train, not only at save. VRAM matched the August 31 capture exactly on both repeats.

ZeRO-3 did not beat the 1-GPU 61.6 GB number on this 7B+checkpointing job. It used both cards. Use it when one 80 GB card is not enough (longer seq, larger batch, 13B / 32B / 70B), or when you already know DDP / ZeRO-2 OOM.

Confirm the checkpoint

ssh -o IdentitiesOnly=yes Ubuntu@YOUR_VM_IP \
  'ls -lh ~/ds-out/0p5-z3/model.safetensors ~/ds-out/7b-z3/model.safetensors ~/ds-out/7b-z3/checkpoint-8/zero_to_fp32.py'

zero_to_fp32.py is the DeepSpeed helper that gathers ZeRO shards into one fp32 folder. This smoke already gathered 16-bit weights via stage3_gather_16bit_weights_on_model_save.