Update Jupyter Notebooks to convert MaxText checkpoint to Hugging Face

src/maxtext/checkpoint_conversion/to_huggingface.py

@@ -366,7 +366,7 @@ def main(argv: Sequence[str]) -> None:
     argv: Command-line arguments, which are parsed by `pyconfig`.
   """
   # Initialize maxtext config
-  config = pyconfig.initialize(argv)
+  config = pyconfig.initialize_pydantic(argv)
   assert (
       config.load_full_state_path == ""
   ), "This script expects parameters, not a full state. Use generate_param_only_checkpoint first if needed."

src/maxtext/examples/rl_llama3_demo.ipynb

@@ -106,6 +106,7 @@
     "from etils import epath\n",
     "import jax\n",
     "\n",
+    "from maxtext.configs import pyconfig\n",
     "from maxtext.trainers.post_train.rl.train_rl import rl_train\n",
     "from maxtext.utils.model_creation_utils import setup_configs_and_devices\n",
     "from maxtext.utils.globals import MAXTEXT_REPO_ROOT, MAXTEXT_PKG_DIR\n",
@@ -292,6 +293,59 @@
     "    sys.exit(1)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Convert MaxText Checkpoint to Hugging Face Format"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config = pyconfig.initialize_pydantic(config_argv)\n",
+    "\n",
+    "# Define the output directory for the Hugging Face checkpoint\n",
+    "hf_output_directory = epath.Path(BASE_OUTPUT_DIRECTORY) / \"hf_checkpoint\"\n",
+    "\n",
+    "# Find the latest MaxText checkpoint\n",
+    "checkpoint_dir = epath.Path(config.checkpoint_dir) / 'actor'\n",
+    "step_dirs = [d.name for d in checkpoint_dir.iterdir() if d.name.isdigit() and d.is_dir()]\n",
+    "if not step_dirs:\n",
+    "    raise ValueError(f\"No checkpoint found in {checkpoint_dir}\")\n",
+    "latest_step = max(step_dirs, key=int)\n",
+    "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"model_params\"\n",
+    "\n",
+    "print(f\"Converting MaxText checkpoint from: {maxtext_checkpoint_path}\")\n",
+    "print(f\"Saving Hugging Face checkpoint to: {hf_output_directory}\")\n",
+    "\n",
+    "# Run the conversion script\n",
+    "env = os.environ.copy()\n",
+    "env[\"JAX_PLATFORMS\"] = \"cpu\"\n",
+    "\n",
+    "subprocess.run(\n",
+    "    [\n",
+    "        sys.executable,\n",
+    "        \"-m\", \"maxtext.checkpoint_conversion.to_huggingface\",\n",
+    "        f\"{MAXTEXT_PKG_DIR}/configs/base.yml\",\n",
+    "        f\"model_name={MODEL_NAME}\",\n",
+    "        f\"load_parameters_path={str(maxtext_checkpoint_path)}\",\n",
+    "        f\"base_output_directory={str(hf_output_directory)}\",\n",
+    "        f\"scan_layers={config.scan_layers}\",\n",
+    "        \"use_multimodal=false\",\n",
+    "        \"skip_jax_distributed_system=True\",\n",
+    "        \"weight_dtype=bfloat16\",\n",
+    "    ],\n",
+    "    check=True,\n",
+    "    env=env\n",
+    ")\n",
+    "\n",
+    "print(\"✓ Conversion completed successfully!\")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

src/maxtext/examples/sft_llama3_demo_tpu.ipynb

@@ -176,7 +176,6 @@
     "MODEL_NAME = \"llama3.1-8b-Instruct\"\n",
     "\n",
     "BASE_OUTPUT_DIRECTORY = f\"{MAXTEXT_PKG_DIR}/sft_llama3_output\"\n",
-    "\n",
     "# set the path to the model checkpoint (including `/0/items`) or leave empty to download from HuggingFace\n",
     "MODEL_CHECKPOINT_PATH = \"\"\n",
     "if not MODEL_CHECKPOINT_PATH:\n",
@@ -322,6 +321,57 @@
     "    sys.exit(1)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Convert MaxText Checkpoint to Hugging Face Format"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define the output directory for the Hugging Face checkpoint\n",
+    "hf_output_directory = epath.Path(BASE_OUTPUT_DIRECTORY) / \"hf_checkpoint\"\n",
+    "\n",
+    "# Find the latest MaxText checkpoint\n",
+    "checkpoint_dir = epath.Path(config.checkpoint_dir)\n",
+    "step_dirs = [d.name for d in checkpoint_dir.iterdir() if d.name.isdigit() and d.is_dir()]\n",
+    "if not step_dirs:\n",
+    "    raise ValueError(f\"No checkpoint found in {checkpoint_dir}\")\n",
+    "latest_step = max(step_dirs, key=int)\n",
+    "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"model_params\"\n",
+    "\n",
+    "print(f\"Converting MaxText checkpoint from: {maxtext_checkpoint_path}\")\n",
+    "print(f\"Saving Hugging Face checkpoint to: {hf_output_directory}\")\n",
+    "\n",
+    "# Run the conversion script\n",
+    "env = os.environ.copy()\n",
+    "env[\"JAX_PLATFORMS\"] = \"cpu\"\n",
+    "\n",
+    "subprocess.run(\n",
+    "    [\n",
+    "        sys.executable,\n",
+    "        \"-m\", \"maxtext.checkpoint_conversion.to_huggingface\",\n",
+    "        f\"{MAXTEXT_PKG_DIR}/configs/base.yml\",\n",
+    "        f\"model_name={MODEL_NAME}\",\n",
+    "        f\"load_parameters_path={str(maxtext_checkpoint_path)}\",\n",
+    "        f\"base_output_directory={str(hf_output_directory)}\",\n",
+    "        f\"scan_layers={config.scan_layers}\",\n",
+    "        \"use_multimodal=false\",\n",
+    "        \"skip_jax_distributed_system=True\",\n",
+    "        \"weight_dtype=bfloat16\",\n",
+    "    ],\n",
+    "    check=True,\n",
+    "    env=env\n",
+    ")\n",
+    "\n",
+    "print(\"✓ Conversion completed successfully!\")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {