[Train] Support FSDP2 checkpoint resume by legitnull · Pull Request #1149 · flagos-ai/FlagScale

legitnull · 2026-03-13T08:41:49Z

PR Category

Train

PR Types

Improvements

PR Description

Refactor to use standard from_pretrained, save_pretrained api
Reorganize qwen_gr00t into proper subpackage (qwen_gr00t/, vlm/qwenvl_backbone)
Add FSDP2 checkpoint resume support (model, optimizer, scheduler, RNG state)
vla serving script now supports arbitrary observation key names

heavyrain-lzy

LGTM

Copilot

Pull request overview

This PR updates FlagScale’s VLA training/serving stack to use a Hugging Face–style checkpoint layout (save_pretrained / from_pretrained), adds FSDP2 resume support (including optimizer/scheduler/RNG), and introduces serve-time preprocessing + observation key remapping for QwenGr00t.

Changes:

Refactor checkpoint save/load to a standard config.json + model.safetensors layout, plus training_state/ for resuming.
Add FSDP2 full-state checkpoint resume utilities and CUDA tests for model/optimizer/training-state round-trips.
Improve QwenGr00t packaging/config dispatch and serving behavior (rename-map + optional serve preprocessor like image resize).

Reviewed changes

Copilot reviewed 30 out of 31 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
tests/unit_tests/train/utils/test_train_utils.py	Adds unit tests for new checkpoint/training-state helpers, including FSDP2 round-trips.
tests/unit_tests/serve/processor/test_image_resize_processor.py	Adds tests for the new serve-time image resize processor step.
tests/unit_tests/models/vla/vlm/test_qwenvl_backbone.py	Updates imports to the renamed QwenVL backbone module.
tests/unit_tests/models/vla/test_qwen_gr00t_config.py	Adds extensive tests for `PreTrainedConfig` dispatch + `QwenGr00tConfig` serialization/validation.
flagscale/train/utils/train_utils.py	Reworks checkpoint saving into `save_checkpoint` + adds FSDP2 load helpers + optimizer/scheduler/rng persistence.
flagscale/train/utils/random_utils.py	Adds safetensors-based RNG state save/load for `random`/`numpy`/`torch`.
flagscale/train/utils/io_utils.py	Adds strict JSON→object deserialization helper (used for scheduler state).
flagscale/train/train_qwen_gr00t.py	Integrates `PreTrainedConfig`, `TrainablePolicy`, FSDP2 optimizer state capture, and resume logic.
flagscale/train/train_config.py	Adds `resume_from`, introduces `train_config.json` save/load helpers, and improves extra-field accessors.
flagscale/train/processor/pipeline.py	Uses `resolve_pretrained_dir` when loading pipeline configs from a directory.
flagscale/serve/run_serve_qwen_gr00t.py	Switches to `TrainablePolicy.from_pretrained`, adds validation, rename-map support, and serve-time processor pipeline.
flagscale/serve/processor/image_resize_processor.py	Introduces an observation image resize processor for serving.
flagscale/serve/processor/init.py	Registers/exports serve processor steps.
flagscale/models/vla/vlm/qwenvl_backbone.py	Refactors VLM config to a dataclass (`QwenVLConfig`) and removes TrainConfig coupling.
flagscale/models/vla/vlm/init.py	Updates exports for renamed backbone module.
flagscale/models/vla/registry.py	Simplifies `build_action_model` API to `**kwargs`.
flagscale/models/vla/qwen_gr00t/task_encoder_qwen_gr00t.py	Adds WDS task encoder under the reorganized qwen_gr00t subpackage.
flagscale/models/vla/qwen_gr00t/modeling_qwen_gr00t.py	Migrates QwenGr00t to `TrainablePolicy` + `PreTrainedConfig`, adds self-contained VLM artifact saving/loading.
flagscale/models/vla/qwen_gr00t/configuration_qwen_gr00t.py	Adds `QwenGr00tConfig` implementing polymorphic config + delta index logic.
flagscale/models/vla/qwen_gr00t/init.py	Exposes QwenGr00t model + config from the new subpackage.
flagscale/models/vla/pretrained_config.py	Introduces polymorphic config base (`PreTrainedConfig`) with registry + serialization.
flagscale/models/vla/base_policy.py	Introduces `TrainablePolicy` base with registry + `save_pretrained`/`from_pretrained`.
flagscale/models/vla/action_model/gr00t_action_header.py	Adds `GR00TActionHeadConfig` + registers `GR00TActionHead`; refactors Flowmatching head init to config-driven.
flagscale/models/vla/action_model/flow_matching.py	Removes old FlowMatchingHead wrapper (replaced by GR00TActionHead).
flagscale/models/vla/action_model/init.py	Updates exports to `GR00TActionHead`.
flagscale/models/vla/init.py	Exports `TrainablePolicy` and removes `FlowMatchingHead` from public API.
flagscale/models/utils/constants.py	Adds `SAFETENSORS_FILE` constant + `resolve_pretrained_dir` helper.
flagscale/inference/inference_qwen_gr00t.py	Switches inference loading to `TrainablePolicy.from_pretrained` + pipeline loading.
examples/qwen_gr00t/conf/train/qwen_gr00t.yaml	Updates config layout (action head type, prompt_template, normalization_mapping, resume_from comment).
examples/qwen_gr00t/conf/serve/qwen_gr00t.yaml	Adds rename-map and serve-time preprocessor (image resize) example.
.github/workflows/unit_tests_common.yml	Installs added deps for unit tests (incl. `datasets`).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

legitnull · 2026-03-25T05:45:31Z

+        # Advance the dataloader iterator to the correct position. The data
+        # ordering is deterministic from the DistributedSampler's own seed
+        # (not the global RNG), so calling next() `step` times moves the
+        # cursor to the right batch. We save/restore the global RNG around
+        # this so the fast-forward's incidental RNG consumption is discarded
+        # and training resumes with the exact RNG state from the checkpoint.
+        saved_rng = serialize_rng_state()
+        # TODO: (yupu) Maybe save/restore the dataloader state?
+        for _ in range(step):
+            next(dl_iter)


todo: will support dataloader resume in the nex PR

Unstaged debug additions and config changes for 0% eval investigation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Remove `torch.device("meta")` context in VLM model loading and the corresponding `to_empty` materialization in `from_pretrained`. Meta device init leaves non-persistent buffers (inv_freq in rotary embeddings) as zeros since they are not saved in safetensors, causing garbage VLM outputs and 0% eval success. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…to qwen_gr00t_fix_lerobot_adapt

…obot_adapt

refactor

7504c21

legitnull force-pushed the qwen_gr00t_fix_lerobot_adapt branch from 8e092a4 to 7504c21 Compare March 13, 2026 08:43

legitnull added 5 commits March 13, 2026 18:06

rename

19227d0

rename

206f37f

resume

57d9314

add ut

db7509c

add more comments

ba832f3

legitnull changed the title ~~Qwen gr00t fix lerobot adapt~~ [Train] Support FSDP2 checkpoint resume Mar 17, 2026

legitnull requested review from 1016qqz, aoyulong and heavyrain-lzy March 17, 2026 09:47

legitnull marked this pull request as ready for review March 17, 2026 09:49

legitnull requested review from cyber-pioneer and zhaoyinglia as code owners March 17, 2026 09:49

legitnull added 3 commits March 19, 2026 10:50

add rename_map in vla serve script

0b84e76

fix ut

17775e3

Merge branch 'main' into qwen_gr00t_fix_lerobot_adapt

6932998

heavyrain-lzy previously approved these changes Mar 19, 2026

View reviewed changes

aoyulong requested a review from Copilot March 19, 2026 04:58

Copilot AI reviewed Mar 19, 2026

View reviewed changes

legitnull and others added 6 commits March 19, 2026 20:30

fix serve

c692c91

WIP: save current state before old-code training experiment

92f9bad

Unstaged debug additions and config changes for 0% eval investigation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

refactor pi

6b271dd

Merge remote-tracking branch 'origin/qwen_gr00t_fix_lerobot_adapt' in…

1499e81

…to qwen_gr00t_fix_lerobot_adapt

Merge remote-tracking branch 'fs_origin/main' into qwen_gr00t_fix_ler…

fe46175

…obot_adapt

legitnull dismissed heavyrain-lzy’s stale review via fe46175 March 23, 2026 07:14

legitnull added 3 commits March 23, 2026 19:38

cotrain

da82732

refactor pi0 serve

651e0f1

add eval

29047de

legitnull added 10 commits March 24, 2026 20:43

test

b217188

test

eec5635

update

10cd84c

fix

0525c75

Merge remote-tracking branch 'fs_origin/main' into qwen_gr00t_fix_ler…

1da6ea2

…obot_adapt

refactor gr00t training

8fdd039

fix serve

a646cb4

fix

f46677f

Merge branch 'main' into qwen_gr00t_fix_lerobot_adapt

2e3395c

fix

7b6a953

Caozhou1995 merged commit cbc199b into flagos-ai:main Mar 25, 2026
4 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Train] Support FSDP2 checkpoint resume#1149

[Train] Support FSDP2 checkpoint resume#1149
Caozhou1995 merged 28 commits into
flagos-ai:mainfrom
legitnull:qwen_gr00t_fix_lerobot_adapt

legitnull commented Mar 13, 2026 •

edited

Loading

Uh oh!

heavyrain-lzy left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

legitnull Mar 25, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

legitnull commented Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Category

PR Types

PR Description

Uh oh!

heavyrain-lzy left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

legitnull Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

legitnull commented Mar 13, 2026 •

edited

Loading