From e08e3d1c2f05b90b7fb8f57614dfbca00117472c Mon Sep 17 00:00:00 2001 From: Sam Anklesaria Date: Mon, 9 Feb 2026 14:45:33 -0600 Subject: [PATCH 1/5] Add 'opt_sharding' option for Variables --- docs_nnx/guides/flax_gspmd.ipynb | 253 ++++++++----------------------- docs_nnx/guides/flax_gspmd.md | 2 +- flax/nnx/training/optimizer.py | 5 +- tests/nnx/spmd_test.py | 22 +++ 4 files changed, 91 insertions(+), 191 deletions(-) diff --git a/docs_nnx/guides/flax_gspmd.ipynb b/docs_nnx/guides/flax_gspmd.ipynb index b1e4c1404..9f0ff9876 100644 --- a/docs_nnx/guides/flax_gspmd.ipynb +++ b/docs_nnx/guides/flax_gspmd.ipynb @@ -2,6 +2,7 @@ "cells": [ { "cell_type": "markdown", + "id": "5d7a3e4e", "metadata": {}, "source": [ "# Scale up on multiple devices\n", @@ -16,8 +17,8 @@ ] }, { - "attachments": {}, "cell_type": "markdown", + "id": "4a50dd2e", "metadata": {}, "source": [ "### Setup" @@ -25,17 +26,10 @@ }, { "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "You have 8 “fake” JAX devices now: [CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]\n" - ] - } - ], + "execution_count": null, + "id": "19242f48", + "metadata": {}, + "outputs": [], "source": [ "from functools import partial\n", "\n", @@ -53,8 +47,8 @@ ] }, { - "attachments": {}, "cell_type": "markdown", + "id": "013b17e5", "metadata": {}, "source": [ "Set up a `2x4` device mesh as the [JAX data sharding tutorial](https://docs.jax.dev/en/latest/sharded-computation.html#key-concept-data-sharding) instructs.\n", @@ -64,7 +58,8 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, + "id": "21cbc359", "metadata": {}, "outputs": [], "source": [ @@ -74,6 +69,7 @@ }, { "cell_type": "markdown", + "id": "b6ba7121", "metadata": {}, "source": [ "> Compatibility Note: This guide covers the [eager sharding feature](https://flax.readthedocs.io/en/latest/flip/4844-var-eager-sharding.html) that greatly simplifies creating sharded model. If your project already used Flax GSPMD API on version `flax<0.12`, you might have turned the feature off to keep your code working. Users can toggle this feature using the `nnx.use_eager_sharding` function." @@ -81,7 +77,8 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, + "id": "608ac912", "metadata": {}, "outputs": [], "source": [ @@ -91,7 +88,7 @@ }, { "cell_type": "markdown", - "id": "c24144d8", + "id": "1145ead5", "metadata": {}, "source": [ "The `nnx.use_eager_sharding` function can also be used as a context manager to toggle the eager sharding feature within a specific scope." @@ -100,7 +97,7 @@ { "cell_type": "code", "execution_count": null, - "id": "2d849e2e", + "id": "9e272034", "metadata": {}, "outputs": [], "source": [ @@ -110,7 +107,7 @@ }, { "cell_type": "markdown", - "id": "c9f808ec", + "id": "61159fa0", "metadata": {}, "source": [ "You can also enable eager sharding on a per-variable basis by passing `eager_sharding=False` during variable initialization. The mesh can also be passed this way." @@ -119,7 +116,7 @@ { "cell_type": "code", "execution_count": null, - "id": "67bbd440", + "id": "e5262d2e", "metadata": {}, "outputs": [], "source": [ @@ -128,63 +125,24 @@ }, { "cell_type": "markdown", + "id": "f36ff900", "metadata": {}, "source": [ "## Shard a single-array model\n", "\n", "Let's begin by sharding the simplest component possible - a Flax variable.\n", "\n", - "When you define a Flax variable, you can pass in a metadata field called `out_sharding`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded.\n", + "When you define a Flax variable, you can pass in a metadata field called `out_sharding`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded. This sharding will also apply to any `nnx.Optimizer` created from the variable. To specify a different sharding for optimization, use the `opt_sharding` metadata field. \n", "\n", "**You must have an existing device mesh** and create a sharding-annotated `nnx.Variable` within its scope. This allows the result variable to be sharded accordingly on those devices. The device mesh can be your actual accelerator mesh, or a dummy fake CPU mesh like in this notebook." ] }, { "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "PartitionSpec(None, 'model')\n" - ] - }, - { - "data": { - "text/html": [ - "
                                                \n",
-       "                                                \n",
-       "                                                \n",
-       "                                                \n",
-       "                                                \n",
-       "  CPU 0,4     CPU 1,5     CPU 2,6     CPU 3,7   \n",
-       "                                                \n",
-       "                                                \n",
-       "                                                \n",
-       "                                                \n",
-       "                                                \n",
-       "
\n" - ], - "text/plain": [ - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0,4\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 1,5\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mCPU 2,6\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mCPU 3,7\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], + "execution_count": null, + "id": "227c2f1f", + "metadata": {}, + "outputs": [], "source": [ "rngs = nnx.Rngs(0)\n", "\n", @@ -199,6 +157,7 @@ }, { "cell_type": "markdown", + "id": "7103496c", "metadata": {}, "source": [ "### Initialize with style\n", @@ -211,6 +170,7 @@ { "cell_type": "code", "execution_count": null, + "id": "aaf5aa43", "metadata": {}, "outputs": [], "source": [ @@ -229,6 +189,7 @@ }, { "cell_type": "markdown", + "id": "55fc12ad", "metadata": {}, "source": [ "### Run the model\n", @@ -242,52 +203,10 @@ }, { "cell_type": "code", - "execution_count": 6, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "PartitionSpec('data', 'model')\n" - ] - }, - { - "data": { - "text/html": [ - "
                                    \n",
-       "                                    \n",
-       "  CPU 0    CPU 1    CPU 2    CPU 3  \n",
-       "                                    \n",
-       "                                    \n",
-       "                                    \n",
-       "                                    \n",
-       "                                    \n",
-       "  CPU 4    CPU 5    CPU 6    CPU 7  \n",
-       "                                    \n",
-       "                                    \n",
-       "                                    \n",
-       "
\n" - ], - "text/plain": [ - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107mCPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82mCPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mCPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\n", - "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;214;97;107m \u001b[0m\u001b[38;2;255;255;255;48;2;140;162;82m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\n", - "\u001b[38;2;0;0;0;48;2;231;203;148m \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m \u001b[0m\n", - "\u001b[38;2;0;0;0;48;2;231;203;148m \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m \u001b[0m\n", - "\u001b[38;2;0;0;0;48;2;231;203;148m \u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148mCPU 4\u001b[0m\u001b[38;2;0;0;0;48;2;231;203;148m \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207mCPU 5\u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148mCPU 6\u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49mCPU 7\u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m \u001b[0m\n", - "\u001b[38;2;0;0;0;48;2;231;203;148m \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m \u001b[0m\n", - "\u001b[38;2;0;0;0;48;2;231;203;148m \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m \u001b[0m\n", - "\u001b[38;2;0;0;0;48;2;231;203;148m \u001b[0m\u001b[38;2;255;255;255;48;2;107;110;207m \u001b[0m\u001b[38;2;255;255;255;48;2;165;81;148m \u001b[0m\u001b[38;2;255;255;255;48;2;140;109;49m \u001b[0m\n" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], + "execution_count": null, + "id": "be77055a", + "metadata": {}, + "outputs": [], "source": [ "# For simple computations, you can get correctly-sharded output without jitting\n", "# In this case, ('data', None) @ (None, 'model') = ('data', 'model')\n", @@ -302,8 +221,8 @@ ] }, { - "attachments": {}, "cell_type": "markdown", + "id": "2c94c22d", "metadata": {}, "source": [ "## Shard a wholesome model\n", @@ -319,7 +238,8 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, + "id": "a9dd8e76", "metadata": {}, "outputs": [], "source": [ @@ -361,6 +281,7 @@ }, { "cell_type": "markdown", + "id": "c145edf1", "metadata": {}, "source": [ "Now a sample training loop, using `jax.jit`." @@ -368,21 +289,10 @@ }, { "cell_type": "code", - "execution_count": 8, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "1.251457\n", - "0.8495563\n", - "0.6590716\n", - "0.5399748\n", - "0.39150265\n" - ] - } - ], + "execution_count": null, + "id": "d5e05ee4", + "metadata": {}, + "outputs": [], "source": [ "@jax.jit\n", "def train_step(model, optimizer, x, y):\n", @@ -411,6 +321,7 @@ }, { "cell_type": "markdown", + "id": "38fb1afe", "metadata": {}, "source": [ "## Profiling\n", @@ -420,17 +331,10 @@ }, { "cell_type": "code", - "execution_count": 9, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "13 ms ± 588 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" - ] - } - ], + "execution_count": null, + "id": "f34fcac4", + "metadata": {}, + "outputs": [], "source": [ "%%timeit\n", "\n", @@ -444,6 +348,7 @@ }, { "cell_type": "markdown", + "id": "c62bb613", "metadata": {}, "source": [ "## Load a sharded model from a checkpoint\n", @@ -455,18 +360,10 @@ }, { "cell_type": "code", - "execution_count": 10, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "PartitionSpec(None, None, 'model')\n", - "(2, 1024, 1024)\n" - ] - } - ], + "execution_count": null, + "id": "eabef455", + "metadata": {}, + "outputs": [], "source": [ "import orbax.checkpoint as ocp\n", "\n", @@ -487,8 +384,8 @@ ] }, { - "attachments": {}, "cell_type": "markdown", + "id": "eeef43b0", "metadata": {}, "source": [ "## Logical axis annotation\n", @@ -500,7 +397,8 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": null, + "id": "c8893f70", "metadata": {}, "outputs": [], "source": [ @@ -544,6 +442,7 @@ }, { "cell_type": "markdown", + "id": "b550fff9", "metadata": {}, "source": [ "If you didn't provide all `sharding_rule` annotations in the model definition, you can apply them at top level by put them into the context via `nnx.logical_axis_rules`." @@ -551,7 +450,8 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": null, + "id": "48e3fc8f", "metadata": {}, "outputs": [], "source": [ @@ -573,8 +473,8 @@ ] }, { - "attachments": {}, "cell_type": "markdown", + "id": "bdb3d4e0", "metadata": {}, "source": [ "### When to use device axis / logical axis\n", @@ -592,6 +492,7 @@ }, { "cell_type": "markdown", + "id": "0735fb3b", "metadata": {}, "source": [ "## Explicit sharding\n", @@ -610,17 +511,9 @@ { "cell_type": "code", "execution_count": null, + "id": "08e6ae32", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "PartitionSpec(None, None, 'model')\n", - "PartitionSpec(None, 'model', None)\n" - ] - } - ], + "outputs": [], "source": [ "# Explicit axis mesh\n", "explicit_mesh = jax.make_mesh((2, 4), ('data', 'model'),\n", @@ -676,6 +569,7 @@ }, { "cell_type": "markdown", + "id": "72b0b797", "metadata": {}, "source": [ "One thing easier in explicit mode is that you can obtain the abstract array tree with shardings via `jax.eval_shape`, instead of calling `nnx.get_abstract_sharding`. This is not possible in auto mode." @@ -683,18 +577,10 @@ }, { "cell_type": "code", - "execution_count": 14, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "PartitionSpec(None, None, 'model')\n", - "PartitionSpec(None, 'model', None)\n" - ] - } - ], + "execution_count": null, + "id": "d4e4052f", + "metadata": {}, + "outputs": [], "source": [ "# Get the sharding tree to load checkpoint with\n", "with jax.set_mesh(explicit_mesh):\n", @@ -706,6 +592,7 @@ }, { "cell_type": "markdown", + "id": "59f1fe99", "metadata": {}, "source": [ "## Further readings\n", @@ -721,20 +608,8 @@ "metadata": { "jupytext": { "formats": "ipynb,md:myst" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.9" } }, "nbformat": 4, - "nbformat_minor": 0 + "nbformat_minor": 5 } diff --git a/docs_nnx/guides/flax_gspmd.md b/docs_nnx/guides/flax_gspmd.md index 90a7fade5..5066306c4 100644 --- a/docs_nnx/guides/flax_gspmd.md +++ b/docs_nnx/guides/flax_gspmd.md @@ -71,7 +71,7 @@ nnx.Param(jnp.ones(4,4), out_sharding=(None, 'model'), eager_sharding=True, mesh Let's begin by sharding the simplest component possible - a Flax variable. -When you define a Flax variable, you can pass in a metadata field called `out_sharding`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded. +When you define a Flax variable, you can pass in a metadata field called `out_sharding`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded. This sharding will also apply to any `nnx.Optimizer` created from the variable. To specify a different sharding for optimization, use the `opt_sharding` metadata field. **You must have an existing device mesh** and create a sharding-annotated `nnx.Variable` within its scope. This allows the result variable to be sharded accordingly on those devices. The device mesh can be your actual accelerator mesh, or a dummy fake CPU mesh like in this notebook. diff --git a/flax/nnx/training/optimizer.py b/flax/nnx/training/optimizer.py index 80ff85b4c..0982a4d2a 100644 --- a/flax/nnx/training/optimizer.py +++ b/flax/nnx/training/optimizer.py @@ -49,7 +49,10 @@ class OptVariable(OptState): def to_opt_state(tree): def _to_opt_state(x): if isinstance(x, Variable): - opt_state = OptVariable(x.get_value(), **x.get_metadata()) # type: ignore + opt_metadata = x.get_metadata() + if 'opt_sharding' in opt_metadata: + opt_metadata['sharding_names'] = opt_metadata['opt_sharding'] + opt_state = OptVariable(x.get_value(), **opt_metadata) # type: ignore else: opt_state = OptArray(x) return opt_state diff --git a/tests/nnx/spmd_test.py b/tests/nnx/spmd_test.py index 8680ae858..0dc2cbbe5 100644 --- a/tests/nnx/spmd_test.py +++ b/tests/nnx/spmd_test.py @@ -90,6 +90,28 @@ def create_module(): assert m.w.shape == (8, 2) assert m.w.sharding.shard_shape(m.w.shape) == (8, 2) + def test_separate_opt_sharding(self): + mesh = jax.make_mesh((2, 2), ("row", "col"), + axis_types=(jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto)) + with jax.set_mesh(mesh): + model = nnx.Linear(4, 2, rngs=nnx.Rngs(0), kernel_metadata={'opt_sharding': ('row', 'col')}) + optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param) + + assert model.kernel.sharding.is_equivalent_to( + NamedSharding(mesh, P()), ndim=2) + assert optimizer.opt_state[0].mu['kernel'].sharding.is_equivalent_to( + NamedSharding(mesh, P('row', 'col')), ndim=2) + + def loss(model, x): + return jnp.sum(model(x)) + + def train_step(model, x): + grads = jax.grad(loss)(model, x) + optimizer.update(model, grads) + + train_step(model, jnp.ones(4)) + train_step(model, jnp.ones(4)) + def test_shard_optimizer_state(self): class Foo(nnx.Module): def __init__(self): From 82a3c1b7bf615e059cf295f6a9247894f44f3d1c Mon Sep 17 00:00:00 2001 From: Sam Anklesaria <61420+samanklesaria@users.noreply.github.com> Date: Mon, 9 Feb 2026 15:00:14 -0600 Subject: [PATCH 2/5] Update flax/nnx/training/optimizer.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- flax/nnx/training/optimizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flax/nnx/training/optimizer.py b/flax/nnx/training/optimizer.py index 0982a4d2a..ad71125dd 100644 --- a/flax/nnx/training/optimizer.py +++ b/flax/nnx/training/optimizer.py @@ -51,7 +51,7 @@ def _to_opt_state(x): if isinstance(x, Variable): opt_metadata = x.get_metadata() if 'opt_sharding' in opt_metadata: - opt_metadata['sharding_names'] = opt_metadata['opt_sharding'] + opt_metadata['sharding_names'] = opt_metadata.pop('opt_sharding') opt_state = OptVariable(x.get_value(), **opt_metadata) # type: ignore else: opt_state = OptArray(x) From 8520c75be577a7904f2a739152757c9f31c51221 Mon Sep 17 00:00:00 2001 From: Sam Anklesaria <61420+samanklesaria@users.noreply.github.com> Date: Thu, 5 Mar 2026 16:31:46 -0600 Subject: [PATCH 3/5] Update out_sharding name Co-authored-by: Cristian Garcia --- flax/nnx/training/optimizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flax/nnx/training/optimizer.py b/flax/nnx/training/optimizer.py index ad71125dd..0797fbf93 100644 --- a/flax/nnx/training/optimizer.py +++ b/flax/nnx/training/optimizer.py @@ -51,7 +51,7 @@ def _to_opt_state(x): if isinstance(x, Variable): opt_metadata = x.get_metadata() if 'opt_sharding' in opt_metadata: - opt_metadata['sharding_names'] = opt_metadata.pop('opt_sharding') + opt_metadata['out_sharding'] = opt_metadata.pop('opt_sharding') opt_state = OptVariable(x.get_value(), **opt_metadata) # type: ignore else: opt_state = OptArray(x) From 0966c0e72a482ae2b96505ea66961ed1c764d6f2 Mon Sep 17 00:00:00 2001 From: Sam Anklesaria Date: Thu, 5 Mar 2026 16:34:04 -0600 Subject: [PATCH 4/5] Fix linter errors --- docs_nnx/guides/flax_gspmd.md | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/docs_nnx/guides/flax_gspmd.md b/docs_nnx/guides/flax_gspmd.md index 5066306c4..44c79adad 100644 --- a/docs_nnx/guides/flax_gspmd.md +++ b/docs_nnx/guides/flax_gspmd.md @@ -22,7 +22,7 @@ If you want to follow explicit sharding style, follow [JAX Explicit Sharding](ht ### Setup -```{code-cell} ipython3 +```{code-cell} from functools import partial import jax @@ -42,28 +42,28 @@ Set up a `2x4` device mesh as the [JAX data sharding tutorial](https://docs.jax. In this guide we use a standard FSDP layout and shard our devices on two axes - `data` and `model`, for doing batch data parallelism and tensor parallelism. -```{code-cell} ipython3 +```{code-cell} # Create an auto-mode mesh of two dimensions and annotate each axis with a name. auto_mesh = jax.make_mesh((2, 4), ('data', 'model')) ``` > Compatibility Note: This guide covers the [eager sharding feature](https://flax.readthedocs.io/en/latest/flip/4844-var-eager-sharding.html) that greatly simplifies creating sharded model. If your project already used Flax GSPMD API on version `flax<0.12`, you might have turned the feature off to keep your code working. Users can toggle this feature using the `nnx.use_eager_sharding` function. -```{code-cell} ipython3 +```{code-cell} nnx.use_eager_sharding(True) assert nnx.using_eager_sharding() ``` The `nnx.use_eager_sharding` function can also be used as a context manager to toggle the eager sharding feature within a specific scope. -```{code-cell} ipython3 +```{code-cell} with nnx.use_eager_sharding(False): assert not nnx.using_eager_sharding() ``` You can also enable eager sharding on a per-variable basis by passing `eager_sharding=False` during variable initialization. The mesh can also be passed this way. -```{code-cell} ipython3 +```{code-cell} nnx.Param(jnp.ones(4,4), out_sharding=(None, 'model'), eager_sharding=True, mesh=auto_mesh) ``` @@ -75,7 +75,7 @@ When you define a Flax variable, you can pass in a metadata field called `out_sh **You must have an existing device mesh** and create a sharding-annotated `nnx.Variable` within its scope. This allows the result variable to be sharded accordingly on those devices. The device mesh can be your actual accelerator mesh, or a dummy fake CPU mesh like in this notebook. -```{code-cell} ipython3 +```{code-cell} rngs = nnx.Rngs(0) with jax.set_mesh(auto_mesh): @@ -93,7 +93,7 @@ When using existing modules, you can apply [`flax.nnx.with_partitioning`](https: Also, you should use `jax.jit` for the whole initialization for maximum performance. This is because without `jax.jit`, a single-device variable must be created first before we apply sharding constraints and then make it sharded, which is wasteful. `jax.jit` will automatically optimize this out. -```{code-cell} ipython3 +```{code-cell} @jax.jit def init_sharded_linear(key): init_fn = nnx.nn.linear.default_kernel_init @@ -115,7 +115,7 @@ You should still make sure to `jax.jit` for maximum performance, and also to exp > Note: You need to `jax.jit` a pure function that takes the model as an argument, instead of jitting the callable model directly. -```{code-cell} ipython3 +```{code-cell} # For simple computations, you can get correctly-sharded output without jitting # In this case, ('data', None) @ (None, 'model') = ('data', 'model') with jax.set_mesh(auto_mesh): @@ -138,7 +138,7 @@ Make note of the following: * [`jax.lax.with_sharding_constraint`](https://docs.jax.dev/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#constraining-shardings-of-intermediates-in-jitted-code): They can help you to enforce specific shardings on intermediate activations. Only works under an auto mode mesh context. -```{code-cell} ipython3 +```{code-cell} class DotReluDot(nnx.Module): def __init__(self, depth: int, rngs: nnx.Rngs): init_fn = nnx.initializers.lecun_normal() @@ -177,7 +177,7 @@ class MultiDotReluDot(nnx.Module): Now a sample training loop, using `jax.jit`. -```{code-cell} ipython3 +```{code-cell} @jax.jit def train_step(model, optimizer, x, y): def loss_fn(model: DotReluDot): @@ -207,7 +207,7 @@ with jax.set_mesh(auto_mesh): If you are using a Google TPU pod or a pod slice, you can create a custom `block_all()` utility function, as defined below, to measure the performance: -```{code-cell} ipython3 +```{code-cell} %%timeit def block_all(xs): @@ -224,7 +224,7 @@ Now you learned how to initialize a sharded model without OOM, but what about sa Make sure you save a model's state, especially if your model shares some variables across modules. Given a You can generate an identical abstract pytree with shardings using Flax’s `nnx.get_abstract_model`. -```{code-cell} ipython3 +```{code-cell} import orbax.checkpoint as ocp # Save the sharded state. @@ -249,7 +249,7 @@ JAX's [automatic](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arr You can provide the mapping along with the annotation as another metadata of the corresponding [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable), or overwrite it at top-level. Check out the `LogicalDotReluDot` example below. -```{code-cell} ipython3 +```{code-cell} # The mapping from alias annotation to the device mesh. sharding_rules = (('batch', 'data'), ('hidden', 'model'), ('embed', None)) @@ -290,7 +290,7 @@ class LogicalMultiDotReluDot(nnx.Module): If you didn't provide all `sharding_rule` annotations in the model definition, you can apply them at top level by put them into the context via `nnx.logical_axis_rules`. -```{code-cell} ipython3 +```{code-cell} with jax.set_mesh(auto_mesh), nnx.logical_axis_rules(sharding_rules): # Model and optimizer logical_model = LogicalMultiDotReluDot(1024, 2, rngs=nnx.Rngs(0)) @@ -334,7 +334,7 @@ For education purposes, we provide a simple Flax model example using explicit sh * Additional dimension from transforms: use `jax.vmap`'s argument `spmd_axis_name`, instead of Flax lifted transforms. -```{code-cell} ipython3 +```{code-cell} # Explicit axis mesh explicit_mesh = jax.make_mesh((2, 4), ('data', 'model'), axis_types=(AxisType.Explicit, AxisType.Explicit)) @@ -389,7 +389,7 @@ assert x.sharding.is_equivalent_to(y.sharding, ndim=2) One thing easier in explicit mode is that you can obtain the abstract array tree with shardings via `jax.eval_shape`, instead of calling `nnx.get_abstract_sharding`. This is not possible in auto mode. -```{code-cell} ipython3 +```{code-cell} # Get the sharding tree to load checkpoint with with jax.set_mesh(explicit_mesh): abs_model = jax.eval_shape( From 61e0d32dc8905196e17b998f014b16c9c977c8e2 Mon Sep 17 00:00:00 2001 From: Sam Anklesaria Date: Thu, 5 Mar 2026 16:44:14 -0600 Subject: [PATCH 5/5] Update name to optimizer_sharding --- docs_nnx/guides/flax_gspmd.ipynb | 2 +- docs_nnx/guides/flax_gspmd.md | 2 +- flax/nnx/training/optimizer.py | 4 ++-- tests/nnx/spmd_test.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs_nnx/guides/flax_gspmd.ipynb b/docs_nnx/guides/flax_gspmd.ipynb index 9f0ff9876..ddae8073a 100644 --- a/docs_nnx/guides/flax_gspmd.ipynb +++ b/docs_nnx/guides/flax_gspmd.ipynb @@ -132,7 +132,7 @@ "\n", "Let's begin by sharding the simplest component possible - a Flax variable.\n", "\n", - "When you define a Flax variable, you can pass in a metadata field called `out_sharding`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded. This sharding will also apply to any `nnx.Optimizer` created from the variable. To specify a different sharding for optimization, use the `opt_sharding` metadata field. \n", + "When you define a Flax variable, you can pass in a metadata field called `optimizer_sharding`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded. This sharding will also apply to any `nnx.Optimizer` created from the variable. To specify a different sharding for optimization, use the `opt_sharding` metadata field. \n", "\n", "**You must have an existing device mesh** and create a sharding-annotated `nnx.Variable` within its scope. This allows the result variable to be sharded accordingly on those devices. The device mesh can be your actual accelerator mesh, or a dummy fake CPU mesh like in this notebook." ] diff --git a/docs_nnx/guides/flax_gspmd.md b/docs_nnx/guides/flax_gspmd.md index 44c79adad..65612b1b5 100644 --- a/docs_nnx/guides/flax_gspmd.md +++ b/docs_nnx/guides/flax_gspmd.md @@ -71,7 +71,7 @@ nnx.Param(jnp.ones(4,4), out_sharding=(None, 'model'), eager_sharding=True, mesh Let's begin by sharding the simplest component possible - a Flax variable. -When you define a Flax variable, you can pass in a metadata field called `out_sharding`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded. This sharding will also apply to any `nnx.Optimizer` created from the variable. To specify a different sharding for optimization, use the `opt_sharding` metadata field. +When you define a Flax variable, you can pass in a metadata field called `optimizer_sharding`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded. This sharding will also apply to any `nnx.Optimizer` created from the variable. To specify a different sharding for optimization, use the `opt_sharding` metadata field. **You must have an existing device mesh** and create a sharding-annotated `nnx.Variable` within its scope. This allows the result variable to be sharded accordingly on those devices. The device mesh can be your actual accelerator mesh, or a dummy fake CPU mesh like in this notebook. diff --git a/flax/nnx/training/optimizer.py b/flax/nnx/training/optimizer.py index 0797fbf93..64313abda 100644 --- a/flax/nnx/training/optimizer.py +++ b/flax/nnx/training/optimizer.py @@ -50,8 +50,8 @@ def to_opt_state(tree): def _to_opt_state(x): if isinstance(x, Variable): opt_metadata = x.get_metadata() - if 'opt_sharding' in opt_metadata: - opt_metadata['out_sharding'] = opt_metadata.pop('opt_sharding') + if 'optimizer_sharding' in opt_metadata: + opt_metadata['out_sharding'] = opt_metadata.pop('optimizer_sharding') opt_state = OptVariable(x.get_value(), **opt_metadata) # type: ignore else: opt_state = OptArray(x) diff --git a/tests/nnx/spmd_test.py b/tests/nnx/spmd_test.py index 0dc2cbbe5..667400d26 100644 --- a/tests/nnx/spmd_test.py +++ b/tests/nnx/spmd_test.py @@ -94,7 +94,7 @@ def test_separate_opt_sharding(self): mesh = jax.make_mesh((2, 2), ("row", "col"), axis_types=(jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto)) with jax.set_mesh(mesh): - model = nnx.Linear(4, 2, rngs=nnx.Rngs(0), kernel_metadata={'opt_sharding': ('row', 'col')}) + model = nnx.Linear(4, 2, rngs=nnx.Rngs(0), kernel_metadata={'optimizer_sharding': ('row', 'col')}) optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param) assert model.kernel.sharding.is_equivalent_to(