Expose gradient accumulation in finetune.sh for low-VRAM fine-tuning#726
Open
maxwbuckley wants to merge 1 commit into
Open
Expose gradient accumulation in finetune.sh for low-VRAM fine-tuning#726maxwbuckley wants to merge 1 commit into
maxwbuckley wants to merge 1 commit into
Conversation
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) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
examples/finetune.shnow forwards aGRADIENT_ACCUMULATION_STEPSenvironment variable (default1— zero behavior change) tolaunch_finetune.py. The--gradient_accumulation_stepsflag was already fully plumbed throughFinetuneConfig→TrainingConfig→TrainingArguments, but unreachable from the documented entry point except via the raw-args escape hatch. Also documents the resulting low-VRAM recipe in the script usage text andgetting_started/hardware_recommendation.md.Why
The single-GPU default (
GLOBAL_BATCH_SIZE=32, no accumulation) sets the fine-tuning floor at 40 GB-class GPUs largely through activation memory. Since activation memory scales with the per-GPU micro-batch while optimization depends only on the effective batch (GLOBAL_BATCH_SIZE x GRADIENT_ACCUMULATION_STEPS), the standard tradeoptimizes with the same effective batch of 32 at a fraction of the activation memory (at some throughput cost), bringing consumer 24–32 GB cards (RTX 4090/5090) into range for the default projector + action-head fine-tune. Weights, gradients, and optimizer state are unaffected.
Verification
finetune.shis executed against a stubbedpythonon PATH and the exact argv asserted — default1, env override forwarded, neighboring args untouched — plusFinetuneConfigvalidation of bad values (0,-1).Gr00tTrainertwice — micro-batch 4 x 2 accumulation vs batch 8 x 1, identical initial weights and sample order — and final parameters asserted equal (allclose, fp32 summation-order tolerance), with a guard that training actually moved the weights. A second test asserts exactly k forwards per optimizer step, so a transformers/trainer change that breaks accumulation semantics fails loudly here.tests/gr00t/experiment/+tests/examples/CPU suites: 35 passed.pre-commitclean.🤖 Generated with Claude Code