Expose gradient accumulation in finetune.sh for low-VRAM fine-tuning

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] \
   [-- <extra launch_finetune.py args>...]
+
+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"

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

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),
+        )

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}"
+    )