add from_pretrained as simple API

Author: A9isha

src/maxtext/configs/post_train/rl.yml

@@ -96,6 +96,7 @@ checkpoint_storage_use_ocdbt: False # For Pathways
 checkpoint_storage_use_zarr3: False # For Pathways
 use_pathways: True
 log_period: 20
+convert_checkpoint_if_possible: True
 
 # ====== Debugging ======
 debug:

src/maxtext/configs/pyconfig.py

@@ -77,16 +77,45 @@ def _module_from_path(path: str) -> str | None:
   return None
 
 
-def _resolve_or_infer_config(argv: list[str]) -> tuple[str, list[str]]:
+def _resolve_or_infer_config(argv: list[str] | None = None, **kwargs) -> tuple[str, list[str]]:
   """Resolves or infers config file path from module."""
+  if argv is None:
+    argv = [""]
+
+  if kwargs.get("base_config"):
+    logger.info("Using config : %s", kwargs["base_config"])
+    return resolve_config_path(kwargs["base_config"]), argv[1:] if len(argv) > 1 else []
+
   if len(argv) >= 2 and argv[1].endswith(".yml"):
     return resolve_config_path(argv[1]), argv[2:]
-  module = _module_from_path(argv[0])
+  module = _module_from_path(argv[0]) if len(argv) > 0 else None
   if module not in _CONFIG_FILE_MAPPING:
-    raise ValueError(f"No config file provided and no default config found for module '{module}'")
-  config_path = os.path.join(MAXTEXT_CONFIGS_DIR, _CONFIG_FILE_MAPPING[module])
-  logger.warning("No config file provided, using default config mapping: %s", config_path)
-  return config_path, argv[1:]
+    config_path = os.path.join(MAXTEXT_CONFIGS_DIR, "base.yml")
+    logger.warning("No config file provided and no default config found for module '%s', using base.yml", module)
+  else:
+    config_path = os.path.join(MAXTEXT_CONFIGS_DIR, _CONFIG_FILE_MAPPING[module])
+    logger.warning("No config file provided, using default config mapping: %s", config_path)
+  remaining_argv = argv[1:] if len(argv) > 1 else []
+
+  return config_path, remaining_argv
+
+
+def _resolve_or_infer_addl_config(**kwargs):
+  """Resolves or infers more configs from module."""
+  inferred_kwargs = {}
+  # if base_output_directory key is not seen
+  if not kwargs.get("base_output_directory"):
+    max_logging.warning("base_output_directory is not provided; Using local directory called maxtext_output")
+    base_output_directory = os.path.abspath("maxtext_output")
+    inferred_kwargs["base_output_directory"] = base_output_directory
+
+  # if hf_access_token key is not seen
+  if not kwargs.get("hf_access_token"):
+    hf_access_token = os.environ.get("HF_TOKEN")
+    if hf_access_token:
+      inferred_kwargs["hf_access_token"] = hf_access_token
+
+  return inferred_kwargs
 
 
 def yaml_key_to_env_key(s: str) -> str:
@@ -280,28 +309,35 @@ def get_keys(self) -> dict[str, Any]:
     return self._flat_config
 
 
-def initialize(argv: list[str], **kwargs) -> HyperParameters:
+def initialize(argv: list[str] | None = None, **kwargs) -> HyperParameters:
   """Initializes the configuration by loading YAML files, and applying CLI, env, and kwarg overrides."""
   pydantic_config = initialize_pydantic(argv, **kwargs)
   config = HyperParameters(pydantic_config)
   return config
 
 
-def initialize_pydantic(argv: list[str], **kwargs) -> MaxTextConfig:
+def initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConfig:
   """Initializes the configuration by loading YAML files, and applying CLI, env, and kwarg overrides.
   Returns pydantic MaxTextConfig 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)
+  config_path, cli_args = _resolve_or_infer_config(argv, **kwargs)
   base_yml_config = _load_config(config_path)
 
   # 2. Get overrides from CLI and kwargs
   cli_cfg = omegaconf.OmegaConf.from_cli(cli_args)
   kwargs_cfg = omegaconf.OmegaConf.create(kwargs)
   overrides_cfg = omegaconf.OmegaConf.merge(cli_cfg, kwargs_cfg)
 
-  # 3. Handle model-specific config
+  temp_cfg1 = omegaconf.OmegaConf.merge(base_yml_config, overrides_cfg)
+  # 3.1. infer more configs if possible
+  temp_cfg1 = _resolve_or_infer_addl_config(**temp_cfg1)
+  # update overrides_cfg with temp_cfg1
+  overrides_cfg = omegaconf.OmegaConf.merge(overrides_cfg, temp_cfg1)
   temp_cfg = omegaconf.OmegaConf.merge(base_yml_config, overrides_cfg)
+
+  # 3.2. Handle model-specific config
+
   model_name = temp_cfg.get("model_name", "default")
   # The architecture for -Instruct v/s base models are the same, so for identifying the
   # architecture we replace "-Instruct" from the model_name and get the base model name
@@ -419,3 +455,13 @@ def initialize_pydantic(argv: list[str], **kwargs) -> MaxTextConfig:
 # Shim for backward compatibility with pyconfig_deprecated_test.py
 validate_and_update_keys = pyconfig_deprecated.validate_and_update_keys
 __all__ = ["initialize", "initialize_pydantic"]
+
+
+class _CallablePyconfigModule(sys.modules[__name__].__class__):
+  """Allows calling the module directly as mt.pyconfig()."""
+
+  def __call__(self, argv: list[str] | None = None, **kwargs) -> HyperParameters:
+    return initialize(argv, **kwargs)
+
+
+sys.modules[__name__].__class__ = _CallablePyconfigModule

src/maxtext/configs/types.py

@@ -1925,6 +1925,11 @@ class DerivedValues(BaseModel):
       None,
       description="The full path to the checkpoint directory, derived from `run_name`.",
   )
+  convert_checkpoint_if_possible: bool = Field(
+      False,
+      description="Whether to convert checkpoint on the fly if not provided via\
+        load_parameters_path or base_output_directory",
+  )
   metrics_dir: None | str = Field(
       None,
       description="The full path to the metrics directory, derived from `run_name`.",

src/maxtext/inference/vllm_decode.py

@@ -241,7 +241,7 @@ def main(argv: Sequence[str]) -> None:
   config = pyconfig.initialize(argv)
 
   if FLAGS.use_tunix:
-    maxtext_model, mesh = model_creation_utils.create_nnx_model(config)
+    maxtext_model, mesh = model_creation_utils.from_pretrained(config)
     decode_with_tunix(config, model=maxtext_model, mesh=mesh)
   else:
     decode_with_vllm(config)

src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py

@@ -251,7 +251,7 @@ def load_weights(self, rng_key: jax.Array) -> None:
       return
 
     with self.mesh, nn.logical_axis_rules(""):
-      model, _ = model_creation_utils.create_nnx_model(
+      model, _ = model_creation_utils.from_pretrained(
           self.maxtext_config, mesh=self.mesh, model_mode=self.model_mode, rng_key=rng_key
       )
       self.model = nnx.data(model)

src/maxtext/trainers/post_train/distillation/train_distill.py

@@ -458,7 +458,7 @@ def get_maxtext_model(config: pyconfig.HyperParameters, mesh: jax.sharding.Mesh)
     The loaded MaxText model.
   """
   max_logging.log(f"Initializing model: {config.model_name}...")
-  model, _ = model_creation_utils.create_nnx_model(config, mesh=mesh)
+  model, _ = model_creation_utils.from_pretrained(config, mesh=mesh)
   return model
 
 

src/maxtext/trainers/post_train/rl/train_rl.py

@@ -47,7 +47,6 @@
 from functools import wraps
 from typing import Sequence
 
-import collections
 import grain
 import jax
 import json
@@ -60,7 +59,6 @@
 from absl import logging as absl_logging
 from etils import epath
 from flax import nnx
-from jax.sharding import Mesh
 from orbax import checkpoint as ocp
 from pprint import pprint
 from transformers import AutoTokenizer
@@ -74,35 +72,10 @@
 
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR
-from maxtext.integration.tunix.tunix_adapter import TunixMaxTextAdapter
 from maxtext.trainers.post_train.rl.evaluate_rl import evaluate
 from maxtext.trainers.post_train.rl import utils_rl
 from maxtext.input_pipeline.instruction_data_processing import load_template_from_file
-from maxtext.utils import max_logging, max_utils, maxtext_utils, model_creation_utils
-
-
-def get_maxtext_model(config, devices=None):
-  """
-  Load MaxText model with Tunix adapter.
-  # Note: pass the path to your scanned checkpoint for 'load_parameters_path'.
-  # To create a scanned checkpoint, you can use /maxtext/src/maxtext/checkpoint_conversion/to_maxtext.py and if
-  # using Pathways, please set `USE_PATHWAYS=1` and use `$((1 - USE_PATHWAYS))` for storage flags:
-  # export USE_PATHWAYS=1
-  # python src/maxtext/checkpoint_conversion/to_maxtext.py \
-  #  --model_name="gemma2-2b" \
-  #  --base_output_directory="/path/to/your/output/directory" \
-  #  --scan_layers=True \
-  #  --checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) \
-  #  --checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS))
-  # Please ensure that you pass the full path ending in `/0/items` for load_parameters_path to train_rl.py i.e.,
-  # load_parameters_path=/path/to/your/output/directory/0/items
-  """
-  model, mesh = model_creation_utils.create_nnx_model(config, devices=devices)
-  with mesh:
-    use_no_op_mappings = "maxtext_config" in config.vllm_additional_config
-    tunix_model = TunixMaxTextAdapter(base_model=model, use_no_op_mappings=use_no_op_mappings)
-    tunix_model.config = None
-  return tunix_model, mesh
+from maxtext.utils import max_logging, max_utils, model_creation_utils
 
 
 def get_dataset(
@@ -150,86 +123,6 @@ def get_dataset(
   return loaded_dataset
 
 
-def setup_configs_and_devices(argv: list[str]):
-  """Setup device allocation and configs for training and inference."""
-  config = pyconfig.initialize_pydantic(argv)
-  devices = jax.devices()
-  if config.num_trainer_slices == -1 and config.num_samplers_slices == -1:
-    max_logging.log("Running RL on a single slice")
-    num_vms = len(devices) // config.chips_per_vm
-    trainer_devices = devices
-    sampler_devices = devices
-    if num_vms >= 2 and config.use_pathways:
-      # Multiple hosts with Pathways - potentially split devices for trainer and sampler
-      # based on trainer_devices_fraction and sampler_devices_fraction
-      max_logging.log(f"{num_vms} VMs detected, allocating trainer and sampler devices, and using Pathways.")
-      num_devices = len(devices)
-      num_trainer_devices = int(num_devices * config.trainer_devices_fraction)
-      num_sampler_devices = int(num_devices * config.sampler_devices_fraction)
-      trainer_devices = devices[:num_trainer_devices]
-      sampler_devices = devices[num_devices - num_sampler_devices :]
-      if config.trainer_devices_fraction != 1.0:
-        max_logging.log(f"Using first {len(trainer_devices)} devices as Trainer devices")
-      if config.sampler_devices_fraction != 1.0:
-        max_logging.log(f"Using last {len(sampler_devices)} devices as Sampler devices")
-    trainer_config = config
-    sampler_config = config
-  elif config.num_trainer_slices > 0 and config.num_samplers_slices > 0:
-    max_logging.log("Running RL with Multislice")
-    devices_by_slice = collections.defaultdict(list)
-    for d in devices:
-      devices_by_slice[d.slice_index].append(d)
-    slice_indices = sorted(devices_by_slice.keys())
-
-    if len(slice_indices) < config.num_trainer_slices + config.num_samplers_slices:
-      raise ValueError("Not enough slices for trainer and samplers")
-
-    trainer_devices = []
-    for i in range(config.num_trainer_slices):
-      trainer_devices.extend(devices_by_slice[slice_indices[i]])
-
-    sampler_devices = []
-    for i in range(config.num_trainer_slices, config.num_trainer_slices + config.num_samplers_slices):
-      sampler_devices.extend(devices_by_slice[slice_indices[i]])
-
-    trainer_devices_per_slice = len(trainer_devices) // config.num_trainer_slices
-    trainer_fsdp = trainer_devices_per_slice
-    tp = config.ici_tensor_parallelism
-    if tp > 1:
-      if trainer_devices_per_slice % tp != 0:
-        raise ValueError(
-            f"trainer_devices_per_slice ({trainer_devices_per_slice}) must be divisible by tensor parallelism ({tp})"
-        )
-      if config.ici_fsdp_parallelism != -1 and config.ici_fsdp_parallelism * tp != trainer_devices_per_slice:
-        raise ValueError(
-            f"ici_fsdp_parallelism ({config.ici_fsdp_parallelism}) * ici_tensor_parallelism ({tp}) must equal "
-            f"devices_per_slice ({trainer_devices_per_slice})"
-        )
-      trainer_fsdp = trainer_devices_per_slice // tp
-
-    trainer_update = {
-        "num_slices": config.num_trainer_slices,
-        "ici_fsdp_parallelism": trainer_fsdp,
-        "ici_tensor_parallelism": tp,
-        "dcn_data_parallelism": config.num_trainer_slices,
-    }
-
-    sampler_update = {
-        "num_slices": config.num_samplers_slices,
-        "ici_fsdp_parallelism": len(sampler_devices) // config.num_samplers_slices,
-        "ici_tensor_parallelism": -1,
-        "dcn_data_parallelism": config.num_samplers_slices,
-    }
-
-    trainer_config = pyconfig.initialize_pydantic(argv, **trainer_update)
-    sampler_config = pyconfig.initialize_pydantic(argv, **sampler_update)
-
-  else:
-    raise ValueError("num_trainer_slices and num_samplers_slices should be both -1 or positive")
-
-  return trainer_config, sampler_config, trainer_devices, sampler_devices
-
-
 def get_rollout_kwargs_for_parallelism(sampler_config, num_sampler_devices):
   """Get rollout kwargs for vLLM rollout when using data parallelism."""
   dp = sampler_config.rollout_data_parallelism
@@ -411,28 +304,6 @@ def _use_raw_prompt(x):
   return train_dataset, test_dataset
 
 
-def create_models_and_meshes(trainer_config, sampler_config, trainer_devices, sampler_devices):
-  """Create reference and actor models and their respective meshes."""
-  max_logging.log("Creating reference model and also meshes for reference and rollout")
-  reference_model, reference_mesh = get_maxtext_model(trainer_config, trainer_devices)
-  devices_array = maxtext_utils.create_device_mesh(sampler_config, sampler_devices)
-  rollout_mesh = Mesh(devices_array, sampler_config.mesh_axes)
-
-  if trainer_config.load_checkpoint_only_once:
-    max_logging.log("Creating policy model by copying reference model instead of restoring from checkpoint again.")
-    with reference_mesh:
-      actor_base_model = nnx.clone(reference_model.base)
-      use_no_op_mappings = "maxtext_config" in trainer_config.vllm_additional_config
-      actor_model = TunixMaxTextAdapter(base_model=actor_base_model, use_no_op_mappings=use_no_op_mappings)
-      actor_model.config = None
-    actor_mesh = reference_mesh
-  else:
-    max_logging.log("Creating policy model with same config as reference model on trainer mesh")
-    actor_model, actor_mesh = get_maxtext_model(trainer_config, trainer_devices)
-
-  return reference_model, reference_mesh, actor_model, actor_mesh, rollout_mesh
-
-
 def create_rl_components(
     trainer_config,
     sampler_config,
@@ -652,7 +523,7 @@ def _reward_fn(**kwargs):
   return rl_cluster, rl_trainer, optimizer
 
 
-def rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices):
+def rl_train(argv: Sequence[str], kwargs: dict):
   """
   Run RL training with the provided configuration.
 
@@ -662,13 +533,21 @@ def rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices):
     trainer_devices: JAX devices for the trainer.
     sampler_devices: JAX devices for the sampler.
   """
+  trainer_config, sampler_config, trainer_devices, sampler_devices = model_creation_utils.setup_configs_and_devices(
+      argv, kwargs
+  )
+
+  reference_model, reference_mesh, actor_model, actor_mesh, rollout_mesh = model_creation_utils.create_models_and_meshes(
+      trainer_config, sampler_config, trainer_devices, sampler_devices
+  )
+
   if not trainer_config.debug.rl:
     # Apply filter to suppress noisy logs
     noise_filter = max_logging.NoisyLogFilter()
     logging.getLogger().addFilter(noise_filter)
     absl_logging.get_absl_logger().addFilter(noise_filter)
+    os.environ["VLLM_LOGGING_LEVEL"] = "ERROR"
 
-  max_logging.log("Starting RL Training")
   if not epath.Path(trainer_config.tensorboard_dir).exists():
     epath.Path(trainer_config.tensorboard_dir).mkdir(parents=True, exist_ok=True)
 
@@ -692,17 +571,14 @@ def rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices):
         break
       pprint(ele)
 
-  reference_model, reference_mesh, actor_model, actor_mesh, rollout_mesh = create_models_and_meshes(
-      trainer_config, sampler_config, trainer_devices, sampler_devices
-  )
-
   if trainer_config.debug.rl:
     max_logging.log("Reference Model initialized successfully")
     nnx.display(reference_model)
     max_logging.log(f"Reference mesh shape: {reference_mesh.shape}")
     max_logging.log("Policy Model initialized successfully")
     nnx.display(actor_model)
     max_logging.log(f"Policy mesh shape: {actor_mesh.shape}")
+    max_logging.log(f"Rollout_mesh shape: {rollout_mesh.shape}")
 
   rl_cluster, rl_trainer, _ = create_rl_components(
       trainer_config,
@@ -759,18 +635,18 @@ def rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices):
   max_logging.warning(f"Post RL Training: {corr=}, {total=}, {accuracy=}%, {partial_accuracy=}%," f" {format_accuracy=}%")
 
 
-def main(argv: Sequence[str]) -> None:
+def main(argv: Sequence[str], kwargs: dict = None) -> None:
   """Main function to run RL training.
 
   Args:
     argv: Command-line arguments.
   """
+  kwargs = kwargs or {}
   pathwaysutils.initialize()
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
 
   max_utils.print_system_information()
-  trainer_config, sampler_config, trainer_devices, sampler_devices = setup_configs_and_devices(argv)
-  rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices)
+  rl_train(argv, kwargs)
 
 
 if __name__ == "__main__":