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