Separate RL configurations into a dedicated config class

Author: SurbhiJainUSC

codecov.yml

@@ -65,4 +65,3 @@ coverage:
         threshold: 10% # fail on 10+ percent degradation
         flags:
           - regular
-

src/maxtext/common/metric_logger.py

@@ -218,7 +218,7 @@ def _log_training_metrics(self, metrics, step):
       moe_lb_loss = scalars.get("learning/moe_lb_loss", 0.0)
       log_parts.append(f"moe_lb_loss: {moe_lb_loss:.3f}")
 
-    if self.config.mtp_num_layers > 0:
+    if getattr(self.config, "mtp_num_layers", 0) > 0:
       mtp_loss = scalars.get("learning/mtp_loss", 0.0)
       log_parts.append(f"main_model_loss: {loss - mtp_loss:.3f}")
       log_parts.append(f"mtp_loss: {mtp_loss:.3f}")
@@ -237,7 +237,7 @@ def _log_eval_metrics(self, metrics, step):
     ]
     if self.config.num_experts > 1:
       log_parts.append(f"avg_moe_lb_loss={scalars['eval/avg_moe_lb_loss']:.3f}")
-    if self.config.mtp_num_layers > 0:
+    if getattr(self.config, "mtp_num_layers", 0) > 0:
       log_parts.extend(
           [
               f"avg_mtp_loss={scalars['eval/avg_mtp_loss']:.3f}",
@@ -258,7 +258,7 @@ def _log_running_eval_metrics(self, metrics, step):
         f"running perplexity={scalars['eval/avg_perplexity']:.3f}",
         f"running total_weights={scalars['eval/total_weights']}",
     ]
-    if self.config.mtp_num_layers > 0:
+    if getattr(self.config, "mtp_num_layers", 0) > 0:
       log_parts.extend(
           [
               f"running mtp_loss={scalars['eval/avg_mtp_loss']:.3f}",

src/maxtext/configs/base.yml

@@ -23,8 +23,6 @@ run_name: ""
 model_name: "default" # override config settings to match a specific model. other than the override, nothing should use this!
 override_model_config: false # When set to true allows overriding model parameters via CLI (or kwargs or env vars) for the purpose of debugging/testing.
 override_logical_axis_rules: false # When set overrides logical axis rules instead of merging them.
-debug:
-  rl: false # RL-specific debugging
 
 normalization_layer_epsilon: 1.e-05 # epsilon value for rmsnorm, layernorm.
 

src/maxtext/configs/inference/vllm.yml

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-base_config: "base.yml"
 attention: "vllm_rpa"
 model_call_mode: "inference"
 
@@ -27,6 +26,11 @@ weight_dtype: bfloat16
 # Allow model config to be overridden by CLI kwargs
 override_model_config: true
 
+# vLLM Adapter Configurations
+vllm_hf_config_path: ""
+vllm_hf_overrides: {}
+vllm_additional_config: {}
+
 
 # -------------- Logical Axis Rules --------------
 mesh_axes: ['data', 'attn_dp', 'model', 'expert', 'attn_dp_expert']

src/maxtext/configs/post_train/rl.yml

@@ -15,8 +15,6 @@
 # RL Configuration
 # This config consolidates common parameters for RL training across different model sizes
 
-base_config: "../base.yml"
-
 # ====== Hardware =====
 trainer_devices_fraction: 0.5
 sampler_devices_fraction: 0.5
@@ -72,7 +70,6 @@ rl:
   # If null, the entire model is resharded at once.
   reshard_chunk_size: null
 
-
 # ====== Models ======
 # for MaxText
 # Model and Tokenizer Configuration
@@ -95,15 +92,16 @@ decoder_layer_input: 'offload'
 query_proj: 'offload'
 key_proj: 'offload'
 value_proj: 'offload'
-checkpoint_storage_use_ocdbt: false # For Pathways
-checkpoint_storage_use_zarr3: false # For Pathways
 use_pathways: true
 log_period: 20
+
+# ====== Checkpointing configs ======
+checkpoint_storage_use_ocdbt: false # For Pathways
+checkpoint_storage_use_zarr3: false # For Pathways
 convert_checkpoint_if_possible: true
 
 # ====== Debugging ======
-debug:
-  rl: true
+debug: false
 # If true, Tunix-managed metrics measurement will be enabled. The metrics will be
 # uploaded to tensorboard.
 enable_tunix_perf_metrics: false
@@ -190,7 +188,6 @@ checkpoint_period: 50
 max_num_checkpoints_to_keep: 10
 
 # ====== Reward ======
-
 reward_exact_answer: 1.0
 reward_white_space_format_match: 1.0
 reward_exact_format_match: 0.1
@@ -199,6 +196,13 @@ reward_ratio_guess_to_answer_high:  0.0
 reward_ratio_guess_to_answer_low: 0.0
 penalty_incorrect_format: 0.0
 penalty_incorrect_answer: 0.0
+# Optional: path to a user-provided Python file with custom reward functions,
+# AND a comma-separated list of function names to import from that file.
+# Each function signature: fn(prompts, completions, tmvp_config, **kwargs) -> list[float].
+# When both are set, the built-in [match_format_exactly, match_format_approximately,
+# check_numbers] reward list is REPLACED entirely.
+reward_functions_path: ''
+reward_functions: ''
 
 # ====== Configuration for math_verify Pool ======
 # Global timeout (seconds) for math_verify calls across all examples in a batch
@@ -235,19 +239,80 @@ eval_dataset_name: 'openai/gsm8k'
 # tmvp_config, x) -> dict with keys {prompts, question, answer}. When empty
 # (default), the built-in utils_rl.process_data is used.
 dataset_processor_path: ''
-# Optional: path to a user-provided Python file with custom reward functions,
-# AND a comma-separated list of function names to import from that file.
-# Each function signature: fn(prompts, completions, tmvp_config, **kwargs) -> list[float].
-# When both are set, the built-in [match_format_exactly, match_format_approximately,
-# check_numbers] reward list is REPLACED entirely.
-reward_functions_path: ''
-reward_functions: ''
 train_split: 'train'
 eval_split: 'test'
-hf_name: 'main' # subset of Hugging Face dataset
-tokenizer_type: 'huggingface'
+hf_subset: 'main' # subset of Hugging Face dataset
 
 ##### MaxText to VLLM Converter validation parameters
 vllm_load_format: dummy # Format to load the model for conversion. Options are "auto", "dummy"
 debug_converter: false # If true, run key coverage check, weight stats, and GCS upload then exit without generation.
 gcs_debug_path: "" # If set and debug_converter=true, upload converted layer-0 and global tensors as .npy to this GCS prefix.
+
+# Logical Axis Rules
+logical_axis_rules: [
+      ['circular_repeats', []],
+
+      # Vocabulary Embedding
+      ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert']],
+      ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'expert']],
+      ['activation_vocab', ['tensor', 'tensor_sequence']],
+      ['activation_vocab', ['tensor']],
+      ['activation_vocab', 'tensor_sequence'],
+      ['vocab', ['tensor', 'tensor_sequence', 'autoregressive']],
+      ['embed_vocab', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
+
+      # Attention
+      ['activation_batch_attn', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
+      ['activation_heads', ['tensor', 'tensor_sequence', 'autoregressive']],
+      ['activation_kv_heads', ['tensor', 'tensor_sequence']],
+      ['activation_length_attn', ['context']],
+      ['activation_q_length', ['context']],
+      ['activation_kv_length', []],
+      ['activation_embed_attn', ['tensor']],
+      ['activation_kv', ['tensor', 'tensor_sequence']],
+      ['activation_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
+      ['activation_kv_head_dim', ['tensor', 'tensor_sequence']],
+      ['heads', ['tensor', 'tensor_sequence', 'autoregressive']],
+      ['q_heads', ['tensor', 'tensor_sequence', 'autoregressive']],
+      ['kv_heads', ['tensor', 'tensor_sequence', 'autoregressive']],
+      ['qkv', []],
+      ['kv', []],
+      ['kv_head_dim', []],
+      ['q_lora', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
+      ['q_lora', ['fsdp', 'context', 'expert']],
+      ["q_lora_up_proj", []],
+      ['kv_lora', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
+      ['kv_lora', ['fsdp', 'context', 'expert']],
+      ["kv_lora_up_proj", []],
+
+      # Mixture of Experts (MoE)
+      ['activation_batch_moe', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
+      ['activation_length_moe', ['context']],
+      ['activation_norm_length_moe', ['tensor_sequence', 'context']],
+      ['activation_embed_moe', ['tensor']],
+      ['activation_mlp_moe', ['tensor', 'tensor_sequence']],
+      ['activation_exp', ['expert']],
+      ['exp', 'expert'],
+      ['mlp_moe', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
+      ['embed_moe', ['fsdp', 'fsdp_transpose', 'context']],
+      ['embed_moe', ['fsdp', 'context']],
+
+      # Standard MLP / Dense Layers / Model Structure
+      ['activation_mlp', ['tensor', 'tensor_sequence']],
+      ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
+      ['activation_length', ['context']],
+      ['activation_norm_length', ['tensor_sequence', 'context']],
+      ['activation_embed', ['tensor']],
+      ['activation_stage', 'stage'],
+      ['mlp', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
+      ['gdn_head', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
+      ['embed', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
+      ['embed', ['fsdp', 'context', 'expert']],
+      ['norm', ['tensor']],
+      ['layers', 'stage'],
+      ['diloco', 'diloco'],
+      ['engram_dim', ['tensor']],
+      ['dense_layers', []],
+      ['moe_layers', []],
+      ['mhc', []],
+  ]

src/maxtext/configs/pyconfig.py

@@ -33,7 +33,6 @@
 from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR, MAXTEXT_ASSETS_ROOT, HF_IDS, MAXTEXT_PKG_DIR
 from maxtext.common.common_types import DecoderBlockType, ShardMode
 from maxtext.configs import types
-from maxtext.configs.types import MaxTextConfig
 from maxtext.inference.inference_utils import str2bool
 from maxtext.utils import max_utils
 from maxtext.utils import max_logging
@@ -238,13 +237,29 @@ def _coerce_to_list(value: Any) -> list[str] | Any:
   return value
 
 
-def _prepare_for_pydantic(raw_keys: dict[str, Any]) -> dict[str, Any]:
+def _check_for_invalid_keys(key: str, valid_fields: set[str]) -> bool:
+  """Recursively checks if the key is valid according to the Pydantic model hierarchy."""
+  if key in valid_fields:
+    return True
+
+  # Check inside nested Pydantic models
+  for field_info in valid_fields.values():
+    annotation = getattr(field_info, "annotation", None)
+    if hasattr(annotation, "model_fields"):
+      if _check_for_invalid_keys(key, annotation.model_fields):
+        return True
+
+  return False
+
+
+def _prepare_for_pydantic(raw_keys: dict[str, Any], config_class: type[Any] = types.MaxTextConfig) -> dict[str, Any]:
   """Prepares the raw dictionary for Pydantic model instantiation."""
   pydantic_kwargs = {}
-  valid_fields = types.MaxTextConfig.model_fields.keys()
+  valid_fields = config_class.model_fields
 
   for key, value in raw_keys.items():
-    if key not in valid_fields:
+    check_flag = _check_for_invalid_keys(key, valid_fields)
+    if not check_flag:
       logger.warning("Ignoring invalid/unsupported field from YAML/CLI: %s", repr(key))
       raise ValueError(f"{key!r} not in {', '.join(map(repr, valid_fields))}.")
 
@@ -318,7 +333,6 @@ def __init__(self, pydantic_config: types.MaxTextConfig):
     final_dict["mu_dtype"] = (
         final_dict["weight_dtype"] if not final_dict["mu_dtype"] else jnp.dtype(final_dict["mu_dtype"])
     )
-
     final_dict["logical_axis_rules"] = _lists_to_tuples(final_dict["logical_axis_rules"])
     final_dict["data_sharding"] = _lists_to_tuples(final_dict["data_sharding"])
 
@@ -379,10 +393,10 @@ def _handle_config_exception(e: Exception):
     raise e
 
 
-def initialize(argv: list[str] | None = None, **kwargs) -> HyperParameters:
+def initialize(argv: list[str] | None = None, config_class: type[Any] = types.MaxTextConfig, **kwargs) -> HyperParameters:
   """Initializes the configuration by loading YAML files, and applying CLI, env, and kwarg overrides."""
   try:
-    pydantic_config = _initialize_pydantic(argv, **kwargs)
+    pydantic_config = _initialize_pydantic(argv, config_class=config_class, **kwargs)
     config = HyperParameters(pydantic_config)
     return config
   except Exception as e:  # pylint: disable=broad-exception-caught
@@ -392,23 +406,23 @@ def initialize(argv: list[str] | None = None, **kwargs) -> HyperParameters:
     raise e
 
 
-def initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConfig:
+def initialize_pydantic(argv: list[str] | None = None, config_class: type[Any] = types.MaxTextConfig, **kwargs) -> Any:
   """Initializes the configuration by loading YAML files, and applying CLI, env, and overrides.
 
-  Returns the pydantic MaxTextConfig class.
+  Returns the pydantic config class.
   """
   try:
-    return _initialize_pydantic(argv, **kwargs)
+    return _initialize_pydantic(argv, config_class=config_class, **kwargs)
   except Exception as e:  # pylint: disable=broad-exception-caught
     if isinstance(e, (SystemExit, KeyboardInterrupt)):
       raise e
     _handle_config_exception(e)
     raise e
 
 
-def _initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConfig:
+def _initialize_pydantic(argv: list[str] | None = None, config_class: type[Any] = types.MaxTextConfig, **kwargs) -> Any:
   """Initializes the configuration by loading YAML files, and applying CLI, env, and kwarg overrides.
-  Returns pydantic MaxTextConfig class whereas `initialize` returns the og `HyperParameters`
+  Returns pydantic config class whereas `initialize` returns the og `HyperParameters`
   """
   # 1. Load base and inherited configs from file(s)
   config_path, cli_args = _resolve_or_infer_config(argv, **kwargs)
@@ -531,7 +545,7 @@ def _initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConf
       except (ValueError, KeyError) as e:
         raise ValueError(f"Couldn't parse value from ENV '{new_proposal}' for key '{k}'") from e
 
-  pydantic_kwargs = _prepare_for_pydantic(raw_keys_dict)
+  pydantic_kwargs = _prepare_for_pydantic(raw_keys_dict, config_class=config_class)
 
   if pydantic_kwargs.get("use_tokamax_splash") and pydantic_kwargs.get("use_jax_splash"):
     raise ValueError("At most one of `use_tokamax_splash` and `use_jax_splash` can be set to True.")
@@ -552,8 +566,8 @@ def _initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConf
 
       compilation_cache.set_cache_dir(os.path.expanduser(pydantic_kwargs["jax_cache_dir"]))
 
-  pydantic_config = types.MaxTextConfig(**pydantic_kwargs)
-  explicit_keys = set((cli_keys | kwargs_keys | env_keys) & frozenset(types.MaxTextConfig.model_fields.keys()))
+  pydantic_config = config_class(**pydantic_kwargs)
+  explicit_keys = set((cli_keys | kwargs_keys | env_keys) & frozenset(config_class.model_fields.keys()))
   object.__setattr__(pydantic_config, "__pydantic_fields_set__", explicit_keys)
   config = HyperParameters(pydantic_config)