Update pyconfig to get config defaults from types.py instead of base.yml

[bvandermoon]

Description

This is Part 2 of the multi-step effort to establish src/maxtext/configs/types.py as the single source of truth for configuration schema and default values in MaxText.

Previously, pyconfig.py required parsing a monolithic 1,400+ line base.yml before passing parameters to Pydantic. This PR updates pyconfig.py so that defaults are provided directly by types.MaxTextConfig, and YAML files are treated as sparse override recipes layered on top of Python defaults.

---

Key Changes

1. Backward-Compatible base.yml Handling:

   • _load_config() now intercepts references to base.yml (e.g. base_config: "base.yml" in existing child YAMLs) and safely treats them as a no-op returning {} with an informational log.
   • Existing workflows, scripts, and child configs referencing base.yml continue running without breaking.

2. Sparse Recipe Loading on types.py Defaults:

   • _initialize_pydantic() no longer assumes a full base YAML dictionary; it merges only the user-provided recipe/model YAML overrides and CLI flags directly into types.MaxTextConfig.
   • Defaults are populated natively by Pydantic rather than read from YAML.

3. Logical Axis Rules Parity:

   • When a recipe YAML does not specify logical_axis_rules, _initialize_pydantic() cleanly falls back to types.DEFAULT_LOGICAL_AXIS_RULES.

4. Dynamic M_* Environment Variable Overrides:

   • Dynamically scans os.environ for any variable starting with M_ (e.g., M_STEPS=100, M_PER_DEVICE_BATCH_SIZE=8, M_GCS_METRICS=True) and maps it to any field in MaxTextConfig.model_fields.

5. Automatic Tokenizer Path Resolution:

   • _prepare_for_pydantic() automatically infers default HuggingFace tokenizer paths when tokenizer_path is omitted in sparse model configs.

6. Unit Tests:

   • Added test_sparse_delta_loading_without_base_yml: Verifies loading sparse YAML recipes directly on top of Pydantic defaults.
   • Added test_dynamic_env_var_overrides: Verifies arbitrary M_<PARAM> environment variable overrides against config fields.

---

Next Steps

• PR 3: Strip base_config: "base.yml" from all 22 child/recipe YAMLs.
• PR 4: Delete base.yml and legacy pyconfig_deprecated.py.

Tests

CI

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request enables sparse configuration loading without a base.yml file and introduces dynamic environment variable overrides prefixed with M_. The reviewer identified several critical issues: a potential silent failure when a specified configuration file is missing, a logic bug in the tokenizer_path fallback for custom models, and a TypeError crash when attempting to override nested configuration fields via environment variables. Expanding the unit tests to cover nested overrides was also recommended.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!