From 0a5d8f49c6a77396eca619734815d00037ecda47 Mon Sep 17 00:00:00 2001 From: haolia Date: Tue, 28 Jul 2026 14:44:51 +0800 Subject: [PATCH] fix(action): make nano DROID policy SFT TOML validate against the schema `examples/toml/sft_config/action_policy_droid_nano.toml` carried `[model.rectified_flow_training_config]` and `[dataloader_train.dataloader]` blocks that `SFTExperimentConfig` (`extra="forbid"`) does not model, so loading it via `--sft-toml` raised `ValidationError` (issue #138). The TOML had never validated against the schema; reference runs used the Hydra `experiment=` path, so the structured-TOML entry point was never exercised. Follow edge practice and keep these values in the registered recipe instead of the TOML: - set the inner dataloader scalars (batch_size/num_workers/prefetch_factor) to the reference values in `action_policy_droid_nano`; - drop both blocks from the TOML (`loss_scale=10` was already set in the recipe). Add a parametrized test that validates every `examples/toml/sft_config/*.toml` against the schema and builds its Hydra overrides. The existing `unittest` job in gpu-tests.yml runs the whole `cosmos_framework/` suite, so the new test is picked up automatically. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../action_policy_droid_nano.py | 6 ++-- .../configs/toml_config/sft_config_test.py | 36 +++++++++++++++++++ .../sft_config/action_policy_droid_nano.toml | 8 ----- 3 files changed, 39 insertions(+), 11 deletions(-) diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py index dfca0b93a..cb86b98e4 100644 --- a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py @@ -172,12 +172,12 @@ tokenizer_spatial_compression_factor=16, tokenizer_temporal_compression_factor=4, dataloader=L(RankPartitionedDataLoader)( - batch_size=1, + batch_size=16, in_order=False, - num_workers=4, + num_workers=16, # host-CPU-bound; lower via CLI on smaller hosts persistent_workers=True, pin_memory=True, - prefetch_factor=4, + prefetch_factor=2, sampler=None, # Shuffling is handled by the dataset (iterable_shuffle=True below): # ActionIterableShuffleDataset streams rank x worker-sharded, episode-order- diff --git a/cosmos_framework/configs/toml_config/sft_config_test.py b/cosmos_framework/configs/toml_config/sft_config_test.py index 3385b41fd..8ff8f3f04 100644 --- a/cosmos_framework/configs/toml_config/sft_config_test.py +++ b/cosmos_framework/configs/toml_config/sft_config_test.py @@ -192,3 +192,39 @@ def test_load_without_custom_section_defaults_empty(self, tmp_path: Path, _dummy config = _load_or_skip(toml_path) assert config.custom == {} + + +# --------------------------------------------------------------------------- # +# 4. Every shipped example TOML validates + builds overrides # +# --------------------------------------------------------------------------- # +_EXAMPLE_TOML_DIR = Path(__file__).parents[3] / "examples" / "toml" / "sft_config" +_EXAMPLE_TOMLS = sorted(_EXAMPLE_TOML_DIR.glob("*.toml")) + + +class TestExampleTomlConfigs: + """Every TOML shipped under ``examples/toml/sft_config/`` must both validate + against ``SFTExperimentConfig`` (``extra="forbid"`` rejects any block the + schema does not register) and produce a Hydra override list. Catches, e.g., a + recipe TOML carrying a ``[model.*]`` / ``[dataloader_train.*]`` sub-block the + schema never modelled, which would raise ``ValidationError`` at load time. + """ + + def test_examples_dir_is_nonempty(self) -> None: + assert _EXAMPLE_TOMLS, f"no example TOMLs found under {_EXAMPLE_TOML_DIR}" + + @pytest.mark.parametrize("toml_path", _EXAMPLE_TOMLS, ids=lambda p: p.name) + def test_example_toml_validates_and_builds_overrides(self, toml_path: Path) -> None: + import tomllib + + raw = tomllib.loads(toml_path.read_text()) + + # 1. Schema validation — extra="forbid" flags un-registered TOML blocks. + try: + SFTExperimentConfig.model_validate(raw) + except ValidationError as exc: # pragma: no cover - failure detail + pytest.fail(f"{toml_path.name} failed schema validation:\n{exc}") + + # 2. Override construction — every leaf must route through PATH_REMAPS. + overrides = build_hydra_overrides(raw) + assert overrides[0] == "--" + assert any(o.startswith("experiment=") for o in overrides), overrides diff --git a/examples/toml/sft_config/action_policy_droid_nano.toml b/examples/toml/sft_config/action_policy_droid_nano.toml index ced8e35f5..92864b06f 100644 --- a/examples/toml/sft_config/action_policy_droid_nano.toml +++ b/examples/toml/sft_config/action_policy_droid_nano.toml @@ -38,9 +38,6 @@ save_ops_regex = ["fmha"] [model.tokenizer] vae_path = "${oc.env:WAN_VAE_PATH}" -[model.rectified_flow_training_config] -loss_scale = 10.0 # generator (diffusion) loss weight, reference value - [optimizer] lr = 2.0e-04 @@ -50,11 +47,6 @@ cycle_lengths = [100000] # long cosine cycle: lr decays slowly across the 1000 [dataloader_train] max_samples_per_batch = 32 # samples packed into each per-rank batch (res480) -[dataloader_train.dataloader] -num_workers = 16 # decode-worker throughput tuning; adjust to host CPU count -batch_size = 16 -prefetch_factor = 2 - [trainer] max_iter = 10000 logging_iter = 50