From 1054b3c9c430769473ff42365ac189de4e4f4304 Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Thu, 16 Jul 2026 12:34:21 +0200 Subject: [PATCH] Expose gradient accumulation in finetune.sh for low-VRAM fine-tuning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit examples/finetune.sh now forwards a GRADIENT_ACCUMULATION_STEPS environment variable (default 1 — no behavior change) to launch_finetune.py, where the --gradient_accumulation_steps flag was already fully plumbed through FinetuneConfig and TrainingConfig but unreachable from the documented entry point without the raw-args escape hatch. This makes the standard low-VRAM trade available as a first-class recipe: GLOBAL_BATCH_SIZE=8 GRADIENT_ACCUMULATION_STEPS=4 optimizes with the same effective batch of 32 as the single-GPU default while activation memory scales down with the micro-batch. Documented in the script usage text and the hardware recommendation guide. Tests cover both halves of the guarantee: - Hermetic plumbing tests execute finetune.sh against a stubbed `python` and assert the exact argv (default 1, env override, neighboring args untouched), plus FinetuneConfig validation of bad values. - Equivalence tests train a deterministic model through the real Gr00tTrainer twice — micro-batch 4 x 2 accumulation vs batch 8 x 1 — and assert parameter-level agreement and one optimizer step per k forwards, so a trainer/transformers change that breaks accumulation semantics fails loudly. Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/finetune.sh | 14 ++ getting_started/hardware_recommendation.md | 1 + tests/examples/test_finetune_sh.py | 132 +++++++++++++ .../experiment/test_gradient_accumulation.py | 184 ++++++++++++++++++ 4 files changed, 331 insertions(+) create mode 100644 tests/examples/test_finetune_sh.py create mode 100644 tests/gr00t/experiment/test_gradient_accumulation.py diff --git a/examples/finetune.sh b/examples/finetune.sh index 8c2dd0396..8fb0bc562 100644 --- a/examples/finetune.sh +++ b/examples/finetune.sh @@ -9,6 +9,11 @@ MAX_STEPS="${MAX_STEPS:-10000}" USE_WANDB="${USE_WANDB:-1}" DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-4}" GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-32}" +# Micro-batches accumulated per optimizer step. The effective batch size is +# GLOBAL_BATCH_SIZE x GRADIENT_ACCUMULATION_STEPS, so lowering GLOBAL_BATCH_SIZE +# while raising this keeps optimization identical at a lower peak VRAM +# (activation memory scales with GLOBAL_BATCH_SIZE / NUM_GPUS). +GRADIENT_ACCUMULATION_STEPS="${GRADIENT_ACCUMULATION_STEPS:-1}" SHARD_SIZE="${SHARD_SIZE:-1024}" NUM_SHARDS_PER_EPOCH="${NUM_SHARDS_PER_EPOCH:-100000}" EPISODE_SAMPLING_RATE="${EPISODE_SAMPLING_RATE:-0.1}" @@ -45,6 +50,14 @@ Usage: bash examples/finetune.sh \ [--save-only-model] \ [--resume-from-checkpoint] \ [-- ...] + +Environment variables: + NUM_GPUS, MASTER_PORT, SAVE_STEPS, MAX_STEPS, USE_WANDB, + DATALOADER_NUM_WORKERS, GLOBAL_BATCH_SIZE, GRADIENT_ACCUMULATION_STEPS, + SHARD_SIZE, NUM_SHARDS_PER_EPOCH, EPISODE_SAMPLING_RATE, DS_WEIGHTS_ALPHA + +Low-VRAM example (same effective batch as the defaults, lower peak memory): + GLOBAL_BATCH_SIZE=8 GRADIENT_ACCUMULATION_STEPS=4 bash examples/finetune.sh ... EOF } @@ -155,6 +168,7 @@ LAUNCH_CMD=( --learning_rate 1e-4 "${WANDB_FLAG[@]}" --global_batch_size "$GLOBAL_BATCH_SIZE" + --gradient_accumulation_steps "$GRADIENT_ACCUMULATION_STEPS" --dataloader_num_workers "$DATALOADER_NUM_WORKERS" --shard_size "$SHARD_SIZE" --num_shards_per_epoch "$NUM_SHARDS_PER_EPOCH" diff --git a/getting_started/hardware_recommendation.md b/getting_started/hardware_recommendation.md index afffbccb7..cd2387b9e 100644 --- a/getting_started/hardware_recommendation.md +++ b/getting_started/hardware_recommendation.md @@ -55,6 +55,7 @@ The table below summarizes end-to-end inference frequency across tested platform - **Default fine-tuning** tunes the projector + diffusion action head (not the full LLM backbone), keeping peak VRAM under ~35 GB per GPU. - **Enabling `--tune-llm` or `--tune-visual`** significantly increases VRAM — 80 GB+ per GPU recommended. - **`--gradient-accumulation-steps`** can compensate for fewer GPUs. For example, 4 GPUs with 8 accumulation steps and per-GPU batch of 8 gives an effective global batch size of 256. +- **Limited VRAM?** Trade micro-batch size for accumulation steps: activation memory scales with the per-GPU micro-batch (`GLOBAL_BATCH_SIZE / NUM_GPUS`), while the optimization itself depends only on the effective batch (`GLOBAL_BATCH_SIZE × GRADIENT_ACCUMULATION_STEPS`). For example, `GLOBAL_BATCH_SIZE=8 GRADIENT_ACCUMULATION_STEPS=4 bash examples/finetune.sh ...` optimizes with the same effective batch of 32 as the single-GPU default, at a fraction of the activation memory (at some throughput cost). Model weights, gradients, and optimizer state are unaffected by this trade. - **Reduce `--num-shards-per-epoch`** if host memory (not VRAM) is limited — this controls how much dataset is preloaded into RAM. --- diff --git a/tests/examples/test_finetune_sh.py b/tests/examples/test_finetune_sh.py new file mode 100644 index 000000000..35644e5b5 --- /dev/null +++ b/tests/examples/test_finetune_sh.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hermetic argument-plumbing tests for examples/finetune.sh. + +The script is executed with a stub ``python`` on PATH that captures the argv it +would have launched, so these tests verify the exact command construction +(single-GPU path) without importing torch or starting a training run. +""" + +from __future__ import annotations + +import os +from pathlib import Path +import subprocess + +import pytest +from test_support.runtime import get_root + + +REPO_ROOT = get_root() +FINETUNE_SH = REPO_ROOT / "examples" / "finetune.sh" + +_STUB_PYTHON = """#!/usr/bin/env bash +printf '%s\\n' "$@" > "$CAPTURE_FILE" +""" + + +def _run_finetune_sh(tmp_path: Path, env_overrides: dict[str, str]) -> list[str]: + """Run finetune.sh with a stubbed ``python`` and return the captured argv.""" + stub_dir = tmp_path / "bin" + stub_dir.mkdir() + stub = stub_dir / "python" + stub.write_text(_STUB_PYTHON) + stub.chmod(0o755) + capture_file = tmp_path / "captured_argv.txt" + + env = os.environ.copy() + env.update(env_overrides) + env["PATH"] = f"{stub_dir}{os.pathsep}{env['PATH']}" + env["CAPTURE_FILE"] = str(capture_file) + env.setdefault("NUM_GPUS", "1") # single-GPU path uses `exec python` + env.setdefault("USE_WANDB", "0") + + result = subprocess.run( + [ + "bash", + str(FINETUNE_SH), + "--base-model-path", + "/fake/model", + "--dataset-path", + "/fake/dataset", + "--embodiment-tag", + "new_embodiment", + "--output-dir", + str(tmp_path / "out"), + ], + env=env, + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, f"finetune.sh failed:\n{result.stderr}" + assert capture_file.exists(), "stub python was never invoked" + return capture_file.read_text().splitlines() + + +def _flag_value(argv: list[str], flag: str) -> str: + """Return the value following ``flag`` in the captured argv.""" + assert flag in argv, f"{flag} not found in argv: {argv}" + return argv[argv.index(flag) + 1] + + +class TestGradientAccumulationPlumbing: + def test_default_is_one(self, tmp_path): + """Without the env var the launcher receives accumulation steps of 1 + (identical behavior to before the flag was exposed).""" + argv = _run_finetune_sh(tmp_path, {}) + assert _flag_value(argv, "--gradient_accumulation_steps") == "1" + + def test_env_var_is_forwarded(self, tmp_path): + argv = _run_finetune_sh(tmp_path, {"GRADIENT_ACCUMULATION_STEPS": "4"}) + assert _flag_value(argv, "--gradient_accumulation_steps") == "4" + + def test_low_vram_recipe_preserves_effective_batch(self, tmp_path): + """The documented low-VRAM recipe: micro-batch 8 x 4 accumulation steps + must reach the launcher exactly as given (effective batch 32, the + single-GPU default).""" + argv = _run_finetune_sh( + tmp_path, + {"GLOBAL_BATCH_SIZE": "8", "GRADIENT_ACCUMULATION_STEPS": "4"}, + ) + global_batch = int(_flag_value(argv, "--global_batch_size")) + accum_steps = int(_flag_value(argv, "--gradient_accumulation_steps")) + assert global_batch == 8 + assert accum_steps == 4 + assert global_batch * accum_steps == 32 + + def test_existing_defaults_unchanged(self, tmp_path): + """Guard: exposing the new env var must not disturb neighboring args.""" + argv = _run_finetune_sh(tmp_path, {}) + assert _flag_value(argv, "--global_batch_size") == "32" + assert _flag_value(argv, "--dataloader_num_workers") == "4" + assert _flag_value(argv, "--embodiment_tag") == "new_embodiment" + + +@pytest.mark.parametrize("bad_value", ["0", "-1"]) +def test_launcher_rejects_invalid_accumulation(bad_value): + """FinetuneConfig validates gradient_accumulation_steps >= 1, so a bad env + var fails fast at config time rather than corrupting the batch math.""" + from gr00t.configs.finetune_config import FinetuneConfig + + with pytest.raises(ValueError, match="gradient_accumulation_steps"): + FinetuneConfig( + base_model_path="/fake/model", + dataset_path="/fake/dataset", + embodiment_tag="new_embodiment", + output_dir="/fake/out", + gradient_accumulation_steps=int(bad_value), + ) diff --git a/tests/gr00t/experiment/test_gradient_accumulation.py b/tests/gr00t/experiment/test_gradient_accumulation.py new file mode 100644 index 000000000..059e29ab0 --- /dev/null +++ b/tests/gr00t/experiment/test_gradient_accumulation.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Gradient-accumulation equivalence tests for Gr00tTrainer. + +The low-VRAM fine-tuning recipe relies on one property: training with +micro-batch ``b`` and ``k`` accumulation steps must perform the same +optimization as micro-batch ``b*k`` with no accumulation. These tests pin that +property on the actual ``Gr00tTrainer`` + ``TrainingArguments`` path (mean-loss +scaling, sequential sample consumption, one optimizer step per ``k`` forwards) +using a deterministic model, so a transformers upgrade or trainer change that +breaks accumulation semantics fails loudly here. + +The GR00T flow-matching loss draws fresh noise per forward call, so full-model +runs are only *statistically* equivalent across batching choices; the trainer +math itself, verified here, is exact. +""" + +from __future__ import annotations + +import copy + +from gr00t.experiment.trainer import Gr00tTrainer +import pytest +import torch +from transformers import PretrainedConfig, PreTrainedModel, TrainingArguments + + +class _TinyRegressionConfig(PretrainedConfig): + model_type = "TinyGradAccumRegression" + + def __init__(self, in_dim: int = 4, **kwargs): + super().__init__(**kwargs) + self.in_dim = in_dim + + +class _TinyRegressionModel(PreTrainedModel): + """Deterministic MSE regressor with Gr00tN1d7's dict-style forward signature.""" + + config_class = _TinyRegressionConfig + + def __init__(self, config: _TinyRegressionConfig): + super().__init__(config) + self.linear = torch.nn.Linear(config.in_dim, 1) + + def forward(self, inputs: dict) -> dict: + pred = self.linear(inputs["x"]).squeeze(-1) + loss = torch.nn.functional.mse_loss(pred, inputs["y"]) + return {"loss": loss} + + +class _FixedDataset(torch.utils.data.Dataset): + """Deterministic samples so both trainer configurations see identical data.""" + + def __init__(self, num_samples: int, in_dim: int): + generator = torch.Generator().manual_seed(1234) + self.x = torch.randn(num_samples, in_dim, generator=generator) + self.y = torch.randn(num_samples, generator=generator) + + def __len__(self) -> int: + return self.x.shape[0] + + def __getitem__(self, idx: int) -> dict: + return {"x": self.x[idx], "y": self.y[idx]} + + +def _collate(examples: list[dict]) -> dict: + """Mirror the production collator shape: a single ``inputs`` dict.""" + return { + "inputs": { + "x": torch.stack([e["x"] for e in examples]), + "y": torch.stack([e["y"] for e in examples]), + } + } + + +def _train( + model: _TinyRegressionModel, + dataset: _FixedDataset, + tmp_path, + *, + per_device_batch: int, + accumulation_steps: int, + max_steps: int, +) -> _TinyRegressionModel: + args = TrainingArguments( + output_dir=str(tmp_path / f"out_b{per_device_batch}_k{accumulation_steps}"), + per_device_train_batch_size=per_device_batch, + gradient_accumulation_steps=accumulation_steps, + max_steps=max_steps, + learning_rate=0.1, + lr_scheduler_type="constant", + seed=0, + report_to=[], + remove_unused_columns=False, + dataloader_num_workers=0, + save_strategy="no", + logging_strategy="no", + disable_tqdm=True, + use_cpu=True, + ) + trainer = Gr00tTrainer( + model=model, + args=args, + train_dataset=dataset, + data_collator=_collate, + ) + trainer.train() + assert trainer.state.global_step == max_steps + return model + + +@pytest.fixture +def dataset() -> _FixedDataset: + return _FixedDataset(num_samples=64, in_dim=4) + + +@pytest.fixture +def initial_model() -> _TinyRegressionModel: + torch.manual_seed(0) + return _TinyRegressionModel(_TinyRegressionConfig()) + + +def test_accumulated_training_matches_full_batch(tmp_path, dataset, initial_model): + """micro-batch 4 x 2 accumulation == batch 8 x 1, parameter for parameter. + + Both runs start from identical weights and consume the same 8 sequential + samples per optimizer step, so the resulting parameters must agree up to + floating-point summation order. + """ + model_full = copy.deepcopy(initial_model) + model_accum = copy.deepcopy(initial_model) + + _train(model_full, dataset, tmp_path, per_device_batch=8, accumulation_steps=1, max_steps=4) + _train(model_accum, dataset, tmp_path, per_device_batch=4, accumulation_steps=2, max_steps=4) + + full_sd = model_full.state_dict() + accum_sd = model_accum.state_dict() + assert full_sd.keys() == accum_sd.keys() + for key in full_sd: + assert torch.allclose(full_sd[key], accum_sd[key], rtol=1e-5, atol=1e-7), ( + f"parameter {key} diverged between accumulated and full-batch training:\n" + f"full={full_sd[key]}\naccum={accum_sd[key]}" + ) + + # Guard against vacuous success: training must actually have moved the weights. + for key, initial_value in initial_model.state_dict().items(): + assert not torch.allclose(full_sd[key], initial_value), ( + f"parameter {key} did not change during training; the equivalence " + f"assertion above proved nothing" + ) + + +def test_accumulation_runs_k_forwards_per_optimizer_step(tmp_path, dataset, initial_model): + """With k accumulation steps the trainer must run k forwards per optimizer + step — i.e. accumulation is actually active, not silently ignored.""" + model = copy.deepcopy(initial_model) + forward_calls = 0 + + def _count_forward(module, args, kwargs): + nonlocal forward_calls + forward_calls += 1 + + handle = model.register_forward_pre_hook(_count_forward, with_kwargs=True) + try: + _train(model, dataset, tmp_path, per_device_batch=4, accumulation_steps=2, max_steps=4) + finally: + handle.remove() + + assert forward_calls == 8, ( + f"expected 4 optimizer steps x 2 accumulation steps = 8 forwards, got {forward_calls}" + )