From 7b033183cf724f8c96014959a379f0588f147727 Mon Sep 17 00:00:00 2001 From: Pian Pawakapan Date: Mon, 17 Aug 2026 15:32:34 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- CONTRIBUTING.md | 7 +- README.md | 2 +- docs/extension.md | 2 + pyproject.toml | 4 +- tests/integration_tests/features.py | 4 +- tests/unit_tests/test_config_manager.py | 15 + tests/unit_tests/test_no_new_cli_options.py | 433 ++++++++++++++++++++ torchtitan/config/README.md | 80 ++++ torchtitan/config/configs.py | 8 + torchtitan_recipes/__init__.py | 7 + torchtitan_recipes/tests.py | 32 ++ 11 files changed, 586 insertions(+), 8 deletions(-) create mode 100644 tests/unit_tests/test_no_new_cli_options.py create mode 100644 torchtitan/config/README.md create mode 100644 torchtitan_recipes/__init__.py create mode 100644 torchtitan_recipes/tests.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 31172cccd0..6b9370e2c3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -50,9 +50,8 @@ Note: To accelerate contributions to and innovations around `torchtitan`, we are - Aim for minimal (if not zero) code change to the model. For the Llama model in `torchtitan`, if one has to make justifiable model change(s): - After the model change, it should still load the original checkpoint correctly. - Document the reasons for the code change, similar to [composability.md](docs/composability.md). -- Keep code modularized, especially for [train.py](torchtitan/train.py), so that it remains easy to copy-paste into a minimal code example. If necessary: - - Introduce new config options/category in [configs.py](torchtitan/config/configs.py). - - Create separate functions/files. +- Keep code modularized, especially for [train.py](torchtitan/train.py), so that it remains easy to copy-paste into a minimal code example. If necessary create separate functions/files. +- The command-line options are frozen: no new `--section.option` flags. A knob that changes the model goes in the model config (dataclass), which is already off the command line. One that belongs to a component goes in that component's config (dataclass), and one with no other owner goes in [configs.py](torchtitan/config/configs.py) after checking with the maintainers. Both need `tyro.conf.Suppress`, since a field there is a command-line option unless you annotate it. See [the configuration doc](torchtitan/config/README.md). ### Proof of Value @@ -75,7 +74,7 @@ When appropriate, one should consider - Adding CPU/GPU unit/integration tests. - To add a unit test, put it in the [tests](tests/) folder and follow the existing test files. - - To add a GPU integration test, create a new `OverrideDefinitions` in [integration_tests](tests/integration_tests/). It will override the default config to run on the Llama 3 debug model (see [config_registry.py](torchtitan/models/llama3/config_registry.py)). + - To add a GPU integration test, add a configuration to [torchtitan_recipes/tests.py](torchtitan_recipes/tests.py) and a new `OverrideDefinitions` naming it in [integration_tests](tests/integration_tests/). Most tests still override a base config from the command line; those are being moved over, so do not add new ones. - Updating [README](README.md) and writing a new note in the [docs](docs/) folder on installation and usage, similar to [float8.md](torchtitan/components/quantization/float8.md). - Following the tensor shape-suffix naming convention for new model code (e.g. `x_BLD`, `q_BLNH`, `out_TNH`), with a per-module legend comment as in [attention.py](torchtitan/models/common/attention.py). Capital suffixes name logical tensor dimensions (not sharding layout) and are scoped per file. - Adding a new file with benchmark results in [benchmarks](benchmarks) folder. diff --git a/README.md b/README.md index bef752dcd1..b94aa84a5b 100644 --- a/README.md +++ b/README.md @@ -75,7 +75,7 @@ We look forward to your contributions! 14. [BF16 optimizer states](docs/bf16_optimizer_states.md) for reduced memory usage 15. Loss, GPU memory, throughput (tokens/sec), TFLOPs, and MFU displayed and logged via [Tensorboard or Weights & Biases](/docs/metrics.md) 16. [Debugging tools](docs/debugging.md) including CPU/GPU profiling, memory profiling, Flight Recorder, etc. -17. All options easily configured via [Python config registry](torchtitan/models/llama3/config_registry.py) with `--module` and `--config` CLI flags +17. All options easily configured in [Python](torchtitan/config/README.md) with `--module` and `--config` CLI flags 18. Structured logging: per-rank trace of key training phases; (see [`torchtitan/observability/structured_logger/README.md`](torchtitan/observability/structured_logger/README.md)) 19. [Helper scripts](scripts/) to - download tokenizers from Hugging Face diff --git a/docs/extension.md b/docs/extension.md index d4107007a9..d4eb1f634b 100644 --- a/docs/extension.md +++ b/docs/extension.md @@ -26,6 +26,8 @@ This is an ongoing effort, and the level of grouping is subject to change. To add custom configuration for an experiment, subclass `Trainer.Config` (or `Trainer` itself) and add new fields. Define config_registry functions that return your custom Config type. +Fields added this way are ordinary command-line options, which is what experiments want. The freeze in [the configuration doc](../torchtitan/config/README.md) applies to core: a field added to a config under `torchtitan/` outside `experiments` needs `tyro.conf.Suppress`. + #### Example To add a custom config section for an experiment: diff --git a/pyproject.toml b/pyproject.toml index c20565e68b..4fce7357cc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,7 +62,9 @@ build-backend = "setuptools.build_meta" [tool.setuptools.packages.find] where = [""] -include = ["torchtitan*"] +# torchtitan_recipes is a second top-level package, listed so that +# tightening the first pattern to "torchtitan.*" cannot silently drop it. +include = ["torchtitan*", "torchtitan_recipes*"] [tool.pytest.ini_options] addopts = ["--showlocals"] # show local variables in tracebacks diff --git a/tests/integration_tests/features.py b/tests/integration_tests/features.py index 7d7ea11ea6..366c51328b 100755 --- a/tests/integration_tests/features.py +++ b/tests/integration_tests/features.py @@ -469,8 +469,8 @@ def build_features_test_list() -> list[OverrideDefinitions]: OverrideDefinitions( [ [ - "--parallelism.data_parallel_shard_degree=2", - "--parallelism.context_parallel_degree=2", + "--module torchtitan_recipes.tests " + "--config llama3_debugmodel_fsdp2_cp2", ] ], "FSDP+CP", diff --git a/tests/unit_tests/test_config_manager.py b/tests/unit_tests/test_config_manager.py index 3b996a6deb..920f791e99 100644 --- a/tests/unit_tests/test_config_manager.py +++ b/tests/unit_tests/test_config_manager.py @@ -59,6 +59,21 @@ def test_missing_both_errors(self): with pytest.raises(ValueError, match="--module is required"): config_manager.parse_args([]) + def test_torchtitan_recipes_package_resolves(self): + """torchtitan_recipes is importable and its configs load.""" + config_manager = ConfigManager() + config = config_manager.parse_args( + [ + "--module", + "torchtitan_recipes.tests", + "--config", + "llama3_debugmodel_fsdp2_cp2", + ] + ) + assert config.model_spec.name == "llama3" + assert config.model_spec.flavor == "debugmodel" + assert config.parallelism.context_parallel_degree == 2 + def test_invalid_model_errors(self): """--module with unknown module name raises ImportError.""" config_manager = ConfigManager() diff --git a/tests/unit_tests/test_no_new_cli_options.py b/tests/unit_tests/test_no_new_cli_options.py new file mode 100644 index 0000000000..6a6b612d84 --- /dev/null +++ b/tests/unit_tests/test_no_new_cli_options.py @@ -0,0 +1,433 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Ensure no more flags are added to CLI""" + +import dataclasses +import importlib +import typing +import unittest +import warnings + +import tyro +from torchtitan.trainer import Trainer + +_FROZEN_CLI_OPTIONS = frozenset( + { + "activation_checkpoint.debug", + "activation_checkpoint.determinism_check", + "activation_checkpoint.force_recompute_mm_shapes_by_fqns", + "activation_checkpoint.memory_budget", + "activation_checkpoint.preserve_rng_state", + "activation_checkpoint.visualize_memory_budget_pareto", + "checkpoint.async_mode", + "checkpoint.create_seed_checkpoint", + "checkpoint.enable", + "checkpoint.enable_first_step_checkpoint", + "checkpoint.exclude_from_loading", + "checkpoint.export_dtype", + "checkpoint.folder", + "checkpoint.initial_load_in_hf", + "checkpoint.initial_load_in_hf_quantized", + "checkpoint.initial_load_model_only", + "checkpoint.initial_load_path", + "checkpoint.interval", + "checkpoint.keep_latest_k", + "checkpoint.last_save_in_hf", + "checkpoint.last_save_model_only", + "checkpoint.load_only", + "checkpoint.load_step", + "comm.init_timeout_seconds", + "comm.mode", + "comm.save_traces_file_prefix", + "comm.save_traces_folder", + "comm.trace_buf_size", + "comm.train_timeout_seconds", + "compile.backend", + "compile.components", + "compile.enable", + "compile.enable_async_tensor_parallel", + "dataloader.build_mrope_positions", + "dataloader.dataset", + "dataloader.dataset_path", + "dataloader.dataset_subset", + "dataloader.generate_timesteps", + "dataloader.image_mean", + "dataloader.image_std", + "dataloader.img_size", + "dataloader.infinite", + "dataloader.load_dataset_kwargs", + "dataloader.max_images_per_batch", + "dataloader.max_patches", + "dataloader.max_patches_per_side", + "dataloader.max_pixels", + "dataloader.min_pixels", + "dataloader.num_workers", + "dataloader.packing_buffer_size", + "dataloader.patch_order", + "dataloader.patch_size", + "dataloader.persistent_workers", + "dataloader.pin_memory", + "dataloader.prefetch_factor", + "dataloader.prompt_dropout_prob", + "dataloader.seed", + "dataloader.sources.dataset", + "dataloader.sources.dataset_path", + "dataloader.sources.infinite", + "dataloader.sources.load_dataset_kwargs", + "dataloader.sources.num_workers", + "dataloader.sources.persistent_workers", + "dataloader.sources.pin_memory", + "dataloader.sources.prefetch_factor", + "dataloader.sources.weight", + "dataloader.spatial_merge_size", + "dataloader.stopping_strategy", + "dataloader.temporal_patch_size", + "dataloader.video_dir", + "dataloader.video_fps", + "dataloader.video_max_frames", + "dataloader.video_min_frames", + "dataloader.weight", + "debug.batch_invariant", + "debug.detect_anomaly", + "debug.deterministic", + "debug.deterministic_warn_only", + "debug.enable_structured_logging", + "debug.moe_force_load_balance", + "debug.print_config", + "debug.save_config_file", + "debug.seed", + "debug.spmd_typechecking", + "dump_folder", + "encoder.autoencoder_path", + "encoder.clip_encoder", + "encoder.random_init", + "encoder.t5_encoder", + "hf_assets_path", + "inference.img_size", + "inference.local_batch_size", + "inference.prompts_path", + "inference.sampling.classifier_free_guidance_scale", + "inference.sampling.denoising_steps", + "inference.sampling.enable_classifier_free_guidance", + "inference.save_img_folder", + "loss.global_vocab_size", + "loss.loss_fn.global_vocab_size", + "loss.loss_fn.mtp_scale", + "loss.mtp_scale", + "loss.num_chunks", + "lr_scheduler.decay_ratio", + "lr_scheduler.decay_type", + "lr_scheduler.min_lr_factor", + "lr_scheduler.total_steps", + "lr_scheduler.warmup_steps", + "metrics.disable_color_printing", + "metrics.enable_tensorboard", + "metrics.enable_wandb", + "metrics.log_freq", + "metrics.save_for_all_ranks", + "metrics.save_tb_folder", + "optimizer.implementation", + "optimizer.optimizer_factory_kwargs_by_name", + "optimizer.param_groups", + "optimizer.param_groups.optimizer_kwargs", + "optimizer.param_groups.optimizer_name", + "optimizer.param_groups.pattern", + "override.imports", + "parallelism.context_parallel_degree", + "parallelism.context_parallel_load_balancer", + "parallelism.context_parallel_ptrr_mask_key", + "parallelism.data_parallel_replicate_degree", + "parallelism.data_parallel_shard_degree", + "parallelism.enable_fsdp_symm_mem", + "parallelism.enable_sequence_parallel", + "parallelism.expert_parallel_degree", + "parallelism.fsdp_reshard_after_forward", + "parallelism.module_fqns_per_model_part", + "parallelism.pipeline_parallel_degree", + "parallelism.pipeline_parallel_first_stage_less_layers", + "parallelism.pipeline_parallel_last_stage_less_layers", + "parallelism.pipeline_parallel_layers_per_stage", + "parallelism.pipeline_parallel_microbatch_size", + "parallelism.pipeline_parallel_schedule", + "parallelism.pipeline_parallel_schedule_csv", + "parallelism.spmd_backend", + "parallelism.tensor_parallel_degree", + "profiler.enable_memory_snapshot", + "profiler.enable_profiling", + "profiler.memory_snapshot_freq", + "profiler.memory_snapshot_max_entries", + "profiler.profile_freq", + "profiler.profiler_active", + "profiler.profiler_repeat", + "profiler.profiler_skip_first", + "profiler.profiler_skip_first_wait", + "profiler.profiler_warmup", + "profiler.save_memory_snapshot_folder", + "profiler.save_traces_folder", + "tokenizer.clip_tokenizer_path", + "tokenizer.image_token", + "tokenizer.max_t5_encoding_len", + "tokenizer.pad_token", + "tokenizer.t5_tokenizer_path", + "tokenizer.test_mode", + "tokenizer.video_token", + "tokenizer.vision_end_token", + "tokenizer.vision_start_token", + "training.disable_cuda_graphs", + "training.dtype", + "training.enable_cpu_offload", + "training.gc_debug", + "training.gc_freq", + "training.global_batch_size", + "training.local_batch_size", + "training.max_norm", + "training.mixed_precision_param", + "training.mixed_precision_reduce", + "training.seq_len", + "training.steps", + "validator.all_timesteps", + "validator.dataloader.build_mrope_positions", + "validator.dataloader.dataset", + "validator.dataloader.dataset_path", + "validator.dataloader.dataset_subset", + "validator.dataloader.generate_timesteps", + "validator.dataloader.image_mean", + "validator.dataloader.image_std", + "validator.dataloader.img_size", + "validator.dataloader.infinite", + "validator.dataloader.load_dataset_kwargs", + "validator.dataloader.max_images_per_batch", + "validator.dataloader.max_patches", + "validator.dataloader.max_patches_per_side", + "validator.dataloader.max_pixels", + "validator.dataloader.min_pixels", + "validator.dataloader.num_workers", + "validator.dataloader.packing_buffer_size", + "validator.dataloader.patch_order", + "validator.dataloader.patch_size", + "validator.dataloader.persistent_workers", + "validator.dataloader.pin_memory", + "validator.dataloader.prefetch_factor", + "validator.dataloader.prompt_dropout_prob", + "validator.dataloader.seed", + "validator.dataloader.sources.dataset", + "validator.dataloader.sources.dataset_path", + "validator.dataloader.sources.infinite", + "validator.dataloader.sources.load_dataset_kwargs", + "validator.dataloader.sources.num_workers", + "validator.dataloader.sources.persistent_workers", + "validator.dataloader.sources.pin_memory", + "validator.dataloader.sources.prefetch_factor", + "validator.dataloader.sources.weight", + "validator.dataloader.spatial_merge_size", + "validator.dataloader.stopping_strategy", + "validator.dataloader.temporal_patch_size", + "validator.dataloader.video_dir", + "validator.dataloader.video_fps", + "validator.dataloader.video_max_frames", + "validator.dataloader.video_min_frames", + "validator.dataloader.weight", + "validator.enable", + "validator.freq", + "validator.sampling.classifier_free_guidance_scale", + "validator.sampling.denoising_steps", + "validator.sampling.enable_classifier_free_guidance", + "validator.save_img_count", + "validator.save_img_folder", + "validator.steps", + } +) + + +def _strip_annotated(field_type): + """The underlying type of an ``Annotated[...]``, or the type itself.""" + return typing.get_args(field_type)[0] if _is_annotated(field_type) else field_type + + +def _is_annotated(field_type) -> bool: + return hasattr(field_type, "__metadata__") + + +def _is_suppressed(field_type) -> bool: + """True when tyro.conf.Suppress hides the field from the command line.""" + return any(m is tyro.conf.Suppress for m in getattr(field_type, "__metadata__", ())) + + +def _cli_options(config, prefix: str = "") -> set[str]: + """Collect the ``section.option`` names tyro exposes for a config. + + Walks the instance. :func:`_declared_cli_options` starts from + ``Trainer.Config`` and never expands subclasses of that root, so a model + that returns its own subclass -- ``FluxTrainer.Config``, with its + ``encoder`` and ``inference`` sections -- is only reachable this way. + """ + options = set() + # Resolved rather than raw: a module using ``from __future__ import + # annotations`` stores its field types as strings, which would hide the + # Suppress annotation. checkpoint.py is one such module. + hints = typing.get_type_hints(type(config), include_extras=True) + for f in dataclasses.fields(config): + field_type = hints.get(f.name, f.type) + if _is_suppressed(field_type): + continue + value = getattr(config, f.name) + name = f"{prefix}{f.name}" + if dataclasses.is_dataclass(value): + options |= _cli_options(value, f"{name}.") + elif value is None and _config_types(field_type): + # A section switched off, such as activation_checkpoint=None. It + # is a subcommand rather than an option, and _declared_cli_options + # already covers what its members expose. + continue + else: + options.add(name) + return options + + +def _subclasses(config_cls: type) -> set[type]: + """``config_cls`` and every imported subclass of it defined in core. + + ``__subclasses__`` sees whatever the process has imported, so an + experiment's config subclass would otherwise appear in the snapshot for + any test run that happened to import it first. The freeze covers core, and + ``torchtitan/experiments`` sets its own rules. + """ + found = {config_cls} + for sub in config_cls.__subclasses__(): + if sub.__module__.startswith("torchtitan.experiments."): + continue + found |= _subclasses(sub) + return found + + +def _config_types(field_type) -> set[type]: + """The config classes a field may hold, unwrapping Annotated and generics. + + Both unions and containers expand. tyro indexes a ``list[ParamGroupConfig]`` + per element, so ``--optimizer.param-groups.0.optimizer-kwargs.lr`` is a real + option; the index is dropped here and the element's fields are recorded once. + + Subclasses expand too, because a field declared as a component base holds + whichever implementation the configuration picked -- ``loss.mtp_scale`` + exists only when the loss is deepseek_v3's. + """ + field_type = _strip_annotated(field_type) + if dataclasses.is_dataclass(field_type): + return _subclasses(field_type) + found: set[type] = set() + for arg in typing.get_args(field_type): + arg = _strip_annotated(arg) + if dataclasses.is_dataclass(arg): + found |= _subclasses(arg) + return found + + +def _declared_cli_options( + config_cls, prefix: str = "", seen: frozenset[type] = frozenset() +) -> set[str]: + """Collect the option names reachable through a config class's annotations. + + Complements :func:`_cli_options`: an instance only shows the components and + union members it happens to hold, so everything else would go unguarded. + + ``seen`` breaks the cycles that subclass expansion creates: + ``ChunkedLossWrapper.Config.loss_fn`` is a ``BaseLoss.Config``, which + expands back to the wrapper. + """ + if config_cls in seen: + return set() + seen = seen | {config_cls} + options = set() + hints = typing.get_type_hints(config_cls, include_extras=True) + for f in dataclasses.fields(config_cls): + field_type = hints.get(f.name, f.type) + if _is_suppressed(field_type): + continue + name = f"{prefix}{f.name}" + nested = _config_types(field_type) + if nested: + for member in nested: + options |= _declared_cli_options(member, f"{name}.", seen) + else: + options.add(name) + return options + + +_GUARDED_CONFIGS = ( + ("llama3", "llama3_debugmodel"), + ("llama3", "sft_debugmodel"), + ("deepseek_v3", "deepseek_v3_debugmodel"), + ("qwen3", "qwen3_debugmodel"), + ("qwen3_5", "qwen35_debugmodel_moe"), + ("gpt_oss", "gpt_oss_debugmodel"), + ("flux", "flux_debugmodel"), + ("kimi_k2_7", "kimi_k2_5_debugmodel"), + ("muse_glimmer", "muse_glimmer_debugmodel_mm"), +) + + +def _guarded_configs(): + """One entry point per model, since each exposes a different surface. + + tyro derives the command line from the selected configuration, so a + component only reachable through one model is only guarded if that model + is walked. A model whose registry needs a dependency this environment + lacks is skipped: the assertion is one-directional, so seeing fewer + options can never fail the test, only cover less. + """ + for model, config_name in _GUARDED_CONFIGS: + try: + registry = importlib.import_module( + f"torchtitan.models.{model}.config_registry" + ) + except ImportError as e: + warnings.warn(f"freeze snapshot skips {model}: {e}", stacklevel=2) + continue + yield getattr(registry, config_name)() + + +class TestCliOptionsFrozen(unittest.TestCase): + def test_every_model_is_guarded(self): + """A new model must be added to _GUARDED_CONFIGS, not silently skipped.""" + from torchtitan.models import _supported_models + + self.assertEqual( + _supported_models - {model for model, _ in _GUARDED_CONFIGS}, + set(), + "These models expose a command-line surface that the freeze does " + "not cover. Add one configuration each to _GUARDED_CONFIGS.", + ) + + def test_no_new_options(self): + # Neither walk alone is the whole surface: an instance shows the + # components it picked, the annotations show everything a declared + # type can hold. Build the configurations first -- the declared walk + # expands subclasses, and a subclass is only visible once the module + # defining it has been imported. + configs = list(_guarded_configs()) + current = _declared_cli_options(Trainer.Config) + for config in configs: + current |= _cli_options(config) + + added = sorted(current - _FROZEN_CLI_OPTIONS) + self.assertFalse( + added, f"The command-line options are frozen, but this adds {added}." + ) + + def test_model_config_tree_is_off_the_cli(self): + """The escape hatch the freeze depends on.""" + hints = typing.get_type_hints(Trainer.Config, include_extras=True) + self.assertTrue( + _is_suppressed(hints["model_spec"]), + "Trainer.Config.model_spec must stay tyro.conf.Suppress: it is " + "what keeps the model config tree off the command line, and " + "therefore what makes the frozen CLI workable.", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/config/README.md b/torchtitan/config/README.md new file mode 100644 index 0000000000..8195e6ec28 --- /dev/null +++ b/torchtitan/config/README.md @@ -0,0 +1,80 @@ +## Configuration + +A run is described by a **full configuration**: a function that returns a complete `Trainer.Config` -- the model, the parallelism degrees, and every optimization choice. Configurations are written in Python, so building a new one is doing configuration programming with TorchTitan components. + +Select one with `--module` (the module that defines the function) and `--config` (the function): + +```bash +NGPU=4 MODULE=torchtitan_recipes.tests CONFIG=llama3_debugmodel_fsdp2_cp2 ./run_train.sh +``` + +The parallelism degrees are in the configuration but the world size is not. So `NGPU` still has to match the product of them. If you want to change any behavior, change the configuration directly -- write your own function, instead of using CLI flags, to return a new `Trainer.Config`: + +```python +# torchtitan_recipes/my_runs.py +def llama3_debugmodel_fsdp2_cp4() -> Trainer.Config: + config = llama3_debugmodel_fsdp2_cp2() + config.parallelism.context_parallel_degree = 4 + config.training.steps = 100 + return config +``` + +The `--section.option` CLI flags still work and still take precedence over the configuration, but only so existing scripts do not break. They are not the way to configure a run any more, and they will be deleted. + +### Where configurations live + +The [torchtitan_recipes](../../torchtitan_recipes/) package holds full configurations -- a recipe is one of these functions. It sits next to `torchtitan` rather than inside it because the two hold different kinds of thing: `torchtitan` ships the model definitions and the classes implementing each optimization; `torchtitan_recipes` only picks combinations of those components. A configuration is also tied to one cluster and one run, so it changes on a different schedule from the library, and shipping one is not the same promise as shipping a class. + +`torchtitan_recipes` is a second top-level package, so an editable install made before it existed does not know about it and `import torchtitan_recipes` fails outside the torchtitan repository root. Re-run `pip install -e .` once if you run outside the repository root. Running from the repository root, as `run_train.sh` and CI do, works either way. + +### Writing your own + +A different cluster usually means a different sharding layout, and therefore a different configuration. That needs no code change: add a function to `torchtitan_recipes`, in a module named for the model, and name it on the command line. (`torchtitan_recipes/tests.py` is separate -- it holds the configurations the integration tests run.) + +```python +# torchtitan_recipes/llama3.py +def llama3_8b_fsdp8_tp2_h200() -> Trainer.Config: + model_spec = model_registry("8B", attn_backend="flex") + return Trainer.Config( + model_spec=model_spec, + parallelism=ParallelismConfig( + data_parallel_shard_degree=8, + tensor_parallel_degree=2, + ), + ... + ) +``` + +`--module` takes any importable module, so a configuration kept outside this repository works the same way: + +```bash +MODULE=my_company_configs.experiments CONFIG=llama3_ablation_7 ./run_train.sh +``` + +### The command-line options are frozen + +The set of `--section.option` CLI flags will not grow. New features express their knobs in the config tree instead, so the way to introduce a new feature is by adding a new configuration, not a new CLI flag. + +Everything already on the command line keeps working, for backward compatibility rather than because it is the recommended path. The eventual goal is to remove the flags entirely and keep only `--module` and `--config`, or even remove tyro completely. + +Frozen means the CLI, not the config dataclasses. New fields still go in the config tree: on the component they belong to, on the model, or -- for the few options with no other home, such as `training.local_batch_size` -- in [configs.py](configs.py), which is not closed, after discussing with the maintainers. + +A field on a component config, or in `configs.py`, needs `tyro.conf.Suppress`: it is a CLI option unless you annotate it, and that annotation is what keeps the CLI from growing while a configuration can still set the field. A field in the model config needs nothing, since `model_spec` is annotated already and takes the whole tree under it off the command line. + +```python +new_job_level_knob: Annotated[int, tyro.conf.Suppress] = 3 +``` + +`Trainer.Config.model_spec` is annotated this way, which is what keeps the whole model config tree off the CLI. + +### What belongs in `torchtitan_recipes` + +What this repository ships, which is deliberately a small set: + +- `tests.py` -- the configurations the integration tests run +- golden configurations verified on specific hardware, named for that hardware so a benchmark run is reproducible from its name alone +- configurations that demonstrate new features + +We do not ship every combination of model, degrees and optimization, because that set is exponential. Your run is your own configuration: add it here without committing it, or keep it in your own package and point `--module` at that. Deriving from a shipped one is a few lines, as above. + +The per-model `config_registry.py` modules, selected with `--module --config `, are the earlier location for the same thing. They keep working and they will eventually be deleted: the model-size baselines they hold, `llama3_8b` and the like, move to `torchtitan_recipes`, so the command line becomes `--module torchtitan_recipes.llama3`. There is no plan for a shim, since a re-export in every model directory would just be a second name for every configuration. diff --git a/torchtitan/config/configs.py b/torchtitan/config/configs.py index bb6c528eab..086afc2c13 100644 --- a/torchtitan/config/configs.py +++ b/torchtitan/config/configs.py @@ -16,6 +16,14 @@ Configs without a clear single owner (or with circular-import constraints) live here. + +Most knobs belong to a component or to the model, not here. But some options +have no suitable home, e.g. ``local_batch_size``, and those can be placed here. +Discuss with the maintainers first if you intend to add one. + +The command-line surface is frozen either way, so annotate a new field with +``tyro.conf.Suppress``, as ``Trainer.Config.model_spec`` does. See +``torchtitan/config/README.md``. """ from dataclasses import dataclass, field diff --git a/torchtitan_recipes/__init__.py b/torchtitan_recipes/__init__.py new file mode 100644 index 0000000000..6594cb86a8 --- /dev/null +++ b/torchtitan_recipes/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""See ``torchtitan/config/README.md``""" diff --git a/torchtitan_recipes/tests.py b/torchtitan_recipes/tests.py new file mode 100644 index 0000000000..9ded50d869 --- /dev/null +++ b/torchtitan_recipes/tests.py @@ -0,0 +1,32 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Full configurations backing the integration tests. + +Each function here is one entry in ``tests/integration_tests``, expressed as +a configuration instead of a base config plus command-line flags. Keeping +them in this package rather than in the test files means CI exercises the +same selection path users do. + +Unrelated to the repository's top-level ``tests/`` package, which holds the +test code itself. +""" + +from torchtitan.models.llama3.config_registry import llama3_debugmodel +from torchtitan.trainer import Trainer + + +def llama3_debugmodel_fsdp2_cp2() -> Trainer.Config: + """Debug model on 4 GPUs: FSDP 2, context parallel 2. + + Derives from ``llama3_debugmodel`` so the two cannot drift, and pins the + parallelism the run needs instead of leaving it to the command line, so + the configuration name is enough to reproduce it. + """ + config = llama3_debugmodel() + config.parallelism.data_parallel_shard_degree = 2 + config.parallelism.context_parallel_degree = 2 + return config