CINXE.COM

<!DOCTYPE html> <html lang="en" data-content_root="" > <head> <meta charset="utf-8" /> <meta name="viewport" content="width=device-width, initial-scale=1.0" /> <title>Scale up Flax Modules on multiple devices</title> <script data-cfasync="false"> document.documentElement.dataset.mode = localStorage.getItem("mode") || ""; document.documentElement.dataset.theme = localStorage.getItem("theme") || ""; </script>  <link href="../../_static/styles/theme.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link href="../../_static/styles/bootstrap.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link href="../../_static/styles/pydata-sphinx-theme.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link href="../../_static/vendor/fontawesome/6.5.2/css/all.min.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="../../_static/vendor/fontawesome/6.5.2/webfonts/fa-solid-900.woff2" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="../../_static/vendor/fontawesome/6.5.2/webfonts/fa-brands-400.woff2" /> <link rel="preload" as="font" type="font/woff2" crossorigin href="../../_static/vendor/fontawesome/6.5.2/webfonts/fa-regular-400.woff2" /> <link rel="stylesheet" type="text/css" href="../../_static/pygments.css" /> <link rel="stylesheet" type="text/css" href="../../_static/styles/sphinx-book-theme.css" /> <link rel="stylesheet" type="text/css" href="../../_static/mystnb.4510f1fc1dee50b3e5859aac5469c37c29e427902b24a333a5f9fcb2f0b3ac41.css" /> <link rel="stylesheet" type="text/css" href="../../_static/sphinx-design.5ea377869091fd0449014c60fc090103.min.css" /> <link rel="stylesheet" type="text/css" href="../../_static/css/flax_theme.css" />  <link rel="preload" as="script" href="../../_static/scripts/bootstrap.js?digest=dfe6caa3a7d634c4db9b" /> <link rel="preload" as="script" href="../../_static/scripts/pydata-sphinx-theme.js?digest=dfe6caa3a7d634c4db9b" /> <script src="../../_static/vendor/fontawesome/6.5.2/js/all.min.js?digest=dfe6caa3a7d634c4db9b"></script> <script data-url_root="../../" id="documentation_options" src="../../_static/documentation_options.js"></script> <script src="../../_static/jquery.js"></script> <script src="../../_static/underscore.js"></script> <script src="../../_static/_sphinx_javascript_frameworks_compat.js"></script> <script src="../../_static/doctools.js"></script> <script src="../../_static/sphinx_highlight.js"></script> <script src="../../_static/scripts/sphinx-book-theme.js"></script> <script src="../../_static/design-tabs.js"></script> <script>DOCUMENTATION_OPTIONS.pagename = 'guides/parallel_training/flax_on_pjit';</script> <link rel="shortcut icon" href="../../_static/flax.png"/> <link rel="index" title="Index" href="../../genindex.html" /> <link rel="search" title="Search" href="../../search.html" /> <link rel="next" title="Model inspection" href="../model_inspection/index.html" /> <link rel="prev" title="Ensembling on multiple devices" href="ensembling.html" /> <meta name="viewport" content="width=device-width, initial-scale=1"/> <meta name="docsearch:language" content="en"/> <script async type="text/javascript" src="/_/static/javascript/readthedocs-addons.js"></script><meta name="readthedocs-project-slug" content="flax-linen" /><meta name="readthedocs-version-slug" content="latest" /><meta name="readthedocs-resolver-filename" content="/guides/parallel_training/flax_on_pjit.html" /><meta name="readthedocs-http-status" content="200" /></head> <body data-bs-spy="scroll" data-bs-target=".bd-toc-nav" data-offset="180" data-bs-root-margin="0px 0px -60%" data-default-mode=""> <div id="pst-skip-link" class="skip-link d-print-none"><a href="#main-content">Skip to main content</a></div> <div id="pst-scroll-pixel-helper"></div> <button type="button" class="btn rounded-pill" id="pst-back-to-top"> Back to top</button> <input type="checkbox" class="sidebar-toggle" id="pst-primary-sidebar-checkbox"/> <label class="overlay overlay-primary" for="pst-primary-sidebar-checkbox"></label> <input type="checkbox" class="sidebar-toggle" id="pst-secondary-sidebar-checkbox"/> <label class="overlay overlay-secondary" for="pst-secondary-sidebar-checkbox"></label> <div class="search-button__wrapper"> <div class="search-button__overlay"></div> <div class="search-button__search-container"> <form class="bd-search d-flex align-items-center" action="../../search.html" method="get"> <input type="search" class="form-control" name="q" id="search-input" placeholder="Search..." aria-label="Search..." autocomplete="off" autocorrect="off" autocapitalize="off" spellcheck="false"/> <kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd>K</kbd> </form></div> </div> <div class="pst-async-banner-revealer d-none"> <aside id="bd-header-version-warning" class="d-none d-print-none" aria-label="Version warning"></aside> </div> <aside class="bd-header-announcement" aria-label="Announcement"> <div class="bd-header-announcement__content"> <a href="https://flax.readthedocs.io/en/latest/index.html" style="text-decoration: none; color: white;" > This site covers the old Flax Linen API. [Explore the new Flax NNX API ✨] </a> </div> </aside> <header class="bd-header navbar navbar-expand-lg bd-navbar d-print-none"> </header> <div class="bd-container"> <div class="bd-container__inner bd-page-width"> <div class="bd-sidebar-primary bd-sidebar"> <div class="sidebar-header-items sidebar-primary__section"> </div> <div class="sidebar-primary-items__start sidebar-primary__section"> <div class="sidebar-primary-item"> <a class="navbar-brand logo" href="../../index.html"> <img src="../../_static/flax.png" class="logo__image only-light" alt=" - Home"/> <script>document.write(`<img src="../../_static/flax.png" class="logo__image only-dark" alt=" - Home"/>`);</script> </a></div> <div class="sidebar-primary-item"> <script> document.write(` <button class="btn search-button-field search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip"> Search <kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd class="kbd-shortcut__modifier">K</kbd> </button> `); </script></div> <div class="sidebar-primary-item"><nav class="bd-links bd-docs-nav" aria-label="Main"> <div class="bd-toc-item navbar-nav active"> <ul class="current nav bd-sidenav"> <li class="toctree-l1"><a class="reference internal" href="../../quick_start.html">Quick start</a></li> <li class="toctree-l1"><a class="reference internal" href="../flax_fundamentals/flax_basics.html">Flax Basics</a></li> <li class="toctree-l1 current active has-children"><a class="reference internal" href="../index.html">Guides</a><details open="open"><summary></summary><ul class="current"> <li class="toctree-l2 has-children"><a class="reference internal" href="../flax_fundamentals/index.html">Flax fundamentals</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference external" href="https://jax.readthedocs.io/en/latest/jax-101/index.html">JAX 101</a></li> <li class="toctree-l3"><a class="reference internal" href="../flax_fundamentals/flax_basics.html">Flax Basics</a></li> <li class="toctree-l3"><a class="reference internal" href="../flax_fundamentals/state_params.html">Managing Parameters and State</a></li> <li class="toctree-l3"><a class="reference internal" href="../flax_fundamentals/setup_or_nncompact.html"><code class="docutils literal notranslate">setup</code> vs <code class="docutils literal notranslate">compact</code></a></li> <li class="toctree-l3"><a class="reference internal" href="../flax_fundamentals/arguments.html">Dealing with Flax Module arguments</a></li> <li class="toctree-l3"><a class="reference internal" href="../flax_fundamentals/rng_guide.html">Randomness and PRNGs in Flax</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../data_preprocessing/index.html">Data preprocessing</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../data_preprocessing/full_eval.html">Processing the entire Dataset</a></li> <li class="toctree-l3"><a class="reference internal" href="../data_preprocessing/loading_datasets.html">Loading datasets</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../training_techniques/index.html">Training techniques</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../training_techniques/batch_norm.html">Batch normalization</a></li> <li class="toctree-l3"><a class="reference internal" href="../training_techniques/dropout.html">Dropout</a></li> <li class="toctree-l3"><a class="reference internal" href="../training_techniques/lr_schedule.html">Learning rate scheduling</a></li> <li class="toctree-l3"><a class="reference internal" href="../training_techniques/transfer_learning.html">Transfer learning</a></li> <li class="toctree-l3"><a class="reference internal" href="../training_techniques/use_checkpointing.html">Save and load checkpoints</a></li> </ul> </details></li> <li class="toctree-l2 current active has-children"><a class="reference internal" href="index.html">Parallel training</a><details open="open"><summary></summary><ul class="current"> <li class="toctree-l3"><a class="reference internal" href="ensembling.html">Ensembling on multiple devices</a></li> <li class="toctree-l3 current active"><a class="current reference internal" href="#">Scale up Flax Modules on multiple devices</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../model_inspection/index.html">Model inspection</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../model_inspection/model_surgery.html">Model surgery</a></li> <li class="toctree-l3"><a class="reference internal" href="../model_inspection/extracting_intermediates.html">Extracting intermediate values</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../converting_and_upgrading/index.html">Converting and upgrading</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../converting_and_upgrading/haiku_migration_guide.html">Migrating from Haiku to Flax</a></li> <li class="toctree-l3"><a class="reference internal" href="../converting_and_upgrading/convert_pytorch_to_flax.html">Convert PyTorch models to Flax</a></li> <li class="toctree-l3"><a class="reference internal" href="../converting_and_upgrading/orbax_upgrade_guide.html">Migrate checkpointing to Orbax</a></li> <li class="toctree-l3"><a class="reference internal" href="../converting_and_upgrading/optax_update_guide.html">Upgrading my codebase to Optax</a></li> <li class="toctree-l3"><a class="reference internal" href="../converting_and_upgrading/linen_upgrade_guide.html">Upgrading my codebase to Linen</a></li> <li class="toctree-l3"><a class="reference internal" href="../converting_and_upgrading/rnncell_upgrade_guide.html">RNNCellBase Upgrade Guide</a></li> <li class="toctree-l3"><a class="reference internal" href="../converting_and_upgrading/regular_dict_upgrade_guide.html">Migrate to regular dicts</a></li> </ul> </details></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../quantization/index.html">Quantization</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../quantization/fp8_basics.html">User Guide on Using FP8</a></li> </ul> </details></li> <li class="toctree-l2"><a class="reference internal" href="../flax_sharp_bits.html">The Sharp Bits</a></li> </ul> </details></li> <li class="toctree-l1 has-children"><a class="reference internal" href="../../examples/index.html">Examples</a><details><summary></summary><ul> <li class="toctree-l2"><a class="reference internal" href="../../examples/core_examples.html">Core examples</a></li> <li class="toctree-l2"><a class="reference internal" href="../../examples/google_research_examples.html">Google Research examples</a></li> <li class="toctree-l2"><a class="reference internal" href="../../examples/repositories_that_use_flax.html">Repositories that use Flax</a></li> <li class="toctree-l2"><a class="reference internal" href="../../examples/community_examples.html">Community examples</a></li> </ul> </details></li> <li class="toctree-l1"><a class="reference internal" href="../../glossary.html">Glossary</a></li> <li class="toctree-l1"><a class="reference internal" href="../../faq.html">Frequently Asked Questions (FAQ)</a></li> <li class="toctree-l1 has-children"><a class="reference internal" href="../../developer_notes/index.html">Developer notes</a><details><summary></summary><ul> <li class="toctree-l2"><a class="reference internal" href="../../developer_notes/module_lifecycle.html">The Flax Module lifecycle</a></li> <li class="toctree-l2"><a class="reference internal" href="../../developer_notes/lift.html">Lifted transformations</a></li> <li class="toctree-l2"><a class="reference external" href="https://github.com/google/flax/tree/main/docs/flip">FLIPs</a></li> </ul> </details></li> <li class="toctree-l1"><a class="reference internal" href="../../philosophy.html">The Flax philosophy</a></li> <li class="toctree-l1"><a class="reference internal" href="../../contributing.html">How to contribute</a></li> <li class="toctree-l1 has-children"><a class="reference internal" href="../../api_reference/index.html">API Reference</a><details><summary></summary><ul> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.config.html">flax.config package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.core.frozen_dict.html">flax.core.frozen_dict package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.cursor.html">flax.cursor package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.errors.html">flax.errors package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.jax_utils.html">flax.jax_utils package</a></li> <li class="toctree-l2 has-children"><a class="reference internal" href="../../api_reference/flax.linen/index.html">flax.linen</a><details><summary></summary><ul> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/module.html">Module</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/init_apply.html">Init/Apply</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/layers.html">Layers</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/activation_functions.html">Activation functions</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/initializers.html">Initializers</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/transformations.html">Transformations</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/inspection.html">Inspection</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/variable.html">Variable dictionary</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/spmd.html">SPMD</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/decorators.html">Decorators</a></li> <li class="toctree-l3"><a class="reference internal" href="../../api_reference/flax.linen/profiling.html">Profiling</a></li> </ul> </details></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.serialization.html">flax.serialization package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.struct.html">flax.struct package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.traceback_util.html">flax.traceback_util package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.training.html">flax.training package</a></li> <li class="toctree-l2"><a class="reference internal" href="../../api_reference/flax.traverse_util.html">flax.traverse_util package</a></li> </ul> </details></li> <li class="toctree-l1"><a class="reference external" href="https://flax.readthedocs.io/en/latest/index.html">Flax NNX</a></li> </ul> </div> </nav></div> </div> <div class="sidebar-primary-items__end sidebar-primary__section"> </div> <div id="rtd-footer-container"></div> </div> <main id="main-content" class="bd-main" role="main"> <div class="sbt-scroll-pixel-helper"></div> <div class="bd-content"> <div class="bd-article-container"> <div class="bd-header-article d-print-none"> <div class="header-article-items header-article__inner"> <div class="header-article-items__start"> <div class="header-article-item"><button class="sidebar-toggle primary-toggle btn btn-sm" title="Toggle primary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip"> </button></div> </div> <div class="header-article-items__end"> <div class="header-article-item"> <div class="article-header-buttons"> <div class="dropdown dropdown-launch-buttons"> <button class="btn dropdown-toggle" type="button" data-bs-toggle="dropdown" aria-expanded="false" aria-label="Launch interactive content"> </button> <ul class="dropdown-menu"> <li><a href="https://colab.research.google.com/github/google/flax/blob/master/docs/guides/parallel_training/flax_on_pjit.ipynb" target="_blank" class="btn btn-sm dropdown-item" title="Launch on Colab" data-bs-placement="left" data-bs-toggle="tooltip" > <img alt="Colab logo" src="../../_static/images/logo_colab.png"> Colab </a> </li> </ul> </div> <a href="https://github.com/google/flax" target="_blank" class="btn btn-sm btn-source-repository-button" title="Source repository" data-bs-placement="bottom" data-bs-toggle="tooltip" > </a> <div class="dropdown dropdown-download-buttons"> <button class="btn dropdown-toggle" type="button" data-bs-toggle="dropdown" aria-expanded="false" aria-label="Download this page"> </button> <ul class="dropdown-menu"> <li><a href="../../_sources/guides/parallel_training/flax_on_pjit.ipynb" target="_blank" class="btn btn-sm btn-download-source-button dropdown-item" title="Download source file" data-bs-placement="left" data-bs-toggle="tooltip" > .ipynb </a> </li> <li> <button onclick="window.print()" class="btn btn-sm btn-download-pdf-button dropdown-item" title="Print to PDF" data-bs-placement="left" data-bs-toggle="tooltip" > .pdf </button> </li> </ul> </div> <button onclick="toggleFullScreen()" class="btn btn-sm btn-fullscreen-button" title="Fullscreen mode" data-bs-placement="bottom" data-bs-toggle="tooltip" > </button> <script> document.write(` <button class="btn btn-sm nav-link pst-navbar-icon theme-switch-button" title="light/dark" aria-label="light/dark" data-bs-placement="bottom" data-bs-toggle="tooltip"> </button> `); </script> <script> document.write(` <button class="btn btn-sm pst-navbar-icon search-button search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip"> </button> `); </script> <button class="sidebar-toggle secondary-toggle btn btn-sm" title="Toggle secondary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip"> </button> </div></div> </div> </div> </div> <div id="jb-print-docs-body" class="onlyprint"> <h1>Scale up Flax Modules on multiple devices</h1>  <div id="print-main-content"> <div id="jb-print-toc"> <div> <h2> Contents </h2> </div> <nav aria-label="Page"> <ul class="visible nav section-nav flex-column"> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#flax-and-jax-jit-scaled-up">Flax and <code class="docutils literal notranslate">jax.jit</code> scaled up</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#setup">Setup</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#define-a-layer">Define a layer</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#define-a-model-with-flax-linen-scan-lifted-transformation">Define a model with <code class="docutils literal notranslate">flax.linen.scan</code> lifted transformation</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#specify-sharding">Specify sharding</a><ul class="nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#the-input-s-sharding">The input’s sharding</a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#the-output-s-sharding">The output’s sharding</a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#compile-the-code">Compile the code</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#inspect-the-module-output">Inspect the Module output</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#compile-the-train-step-and-inference">Compile the train step and inference</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#profiling">Profiling</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#logical-axis-annotation">Logical axis annotation</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#when-to-use-device-axis-logical-axis">When to use device axis / logical axis</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#save-the-data">Save the data</a></li> </ul> </nav> </div> </div> </div> <div id="searchbox"></div> <article class="bd-article"> <div class="tex2jax_ignore mathjax_ignore section" id="scale-up-flax-modules-on-multiple-devices"> <h1>Scale up Flax Modules on multiple devices<a class="headerlink" href="#scale-up-flax-modules-on-multiple-devices" title="Permalink to this heading">#</a></h1> This guide shows how to scale up <a class="reference external" href="https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html">Flax Modules</a> on multiple devices and hosts using <a class="reference external" href="https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html"><code class="docutils literal notranslate">jax.jit</code></a> (formerly <a class="reference external" href="https://jax.readthedocs.io/en/latest/jax.experimental.pjit.html#module-jax.experimental.pjit"><code class="docutils literal notranslate">experimental.pjit</code></a>) and <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/index.html"><code class="docutils literal notranslate">flax.linen</code></a>. <div class="section" id="flax-and-jax-jit-scaled-up"> <h2>Flax and <code class="docutils literal notranslate">jax.jit</code> scaled up<a class="headerlink" href="#flax-and-jax-jit-scaled-up" title="Permalink to this heading">#</a></h2> <a class="reference external" href="https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html"><code class="docutils literal notranslate">jax.jit</code></a> follows the <a class="reference external" href="https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD">Single Program Multi Data (SPMD)</a> paradigm and automatically compiles your code to run it on multiple devices. You need to only specify how you want the input and output of your code to be partitioned, and the compiler will figure out how to: 1) partition everything inside; and 2) compile inter-device communications. Flax provides several functionalities that can help you use auto-SPMD on <a class="reference external" href="https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html">Flax Modules</a>, including: <ol class="arabic simple"> <li>An interface to specify partitions of your data when defining <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html"><code class="docutils literal notranslate">flax.linen.Module</code></a>.</li> <li>Utility functions to generate the sharding information that <code class="docutils literal notranslate">jax.jit</code> requires to run.</li> <li>An interface to customize your axis names called “logical axis annotations” to decouple both your Module code and partition plan to experiment with different partition layouts more easily.</li> </ol> You can learn more about <code class="docutils literal notranslate">jax.jit</code> APIs for scaling up in <a class="reference external" href="https://jax.readthedocs.io/en/latest/multi_process.html">JAX in multi-process environments</a> and <a class="reference external" href="https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html">Distributed arrays and automatic parallelization</a> on JAX’s documentation site. </div> <div class="section" id="setup"> <h2>Setup<a class="headerlink" href="#setup" title="Permalink to this heading">#</a></h2> Import some necessary dependencies. Note: This guide uses the <code class="docutils literal notranslate">--xla_force_host_platform_device_count=8</code> flag to emulate multiple devices in a CPU environment in a Google Colab/Jupyter Notebook. You don’t need this if you are already using a multi-device TPU environment. <div class="cell tag_skip-execution docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre># Once Flax v0.6.10 is released, there is no need to do this. # ! pip3 install -qq "git+https://github.com/google/flax.git@main#egg=flax" </pre></div> </div> </div> </div> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>import os os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' </pre></div> </div> </div> </div> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>import functools from typing import Optional, Callable import numpy as np import jax from jax import lax, random, numpy as jnp import flax from flax import struct, traverse_util, linen as nn from flax.core import freeze, unfreeze from flax.training import train_state, checkpoints import optax # Optax for common losses and optimizers. </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output stderr highlight-myst-ansi notranslate"><div class="highlight"><pre>2024-11-20 21:52:15.791706: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered WARNING: All log messages before absl::InitializeLog() is called are written to STDERR E0000 00:00:1732139535.811354 945 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered E0000 00:00:1732139535.817305 945 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered </pre></div> </div> </div> </div> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>print(f'We have 8 fake JAX devices now: {jax.devices()}') </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre>We 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)] </pre></div> </div> </div> </div> The code below shows how to import and set up the JAX-level device API, following JAX’s <a class="reference external" href="https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html">Distributed arrays and automatic parallelization</a> guide: <ol class="arabic simple"> <li>Start a 2x4 device <code class="docutils literal notranslate">mesh</code> (8 devices) using JAX’s <code class="docutils literal notranslate">mesh_utils.create_device_mesh</code>. This layout is the same as the one of a <a class="reference external" href="https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#single_tpu_board">TPU v3-8</a>.</li> <li>Annotate each axis with a name using the <code class="docutils literal notranslate">axis_names</code> parameter in <code class="docutils literal notranslate">jax.sharding.Mesh</code>. A typical way to annotate axis names is <code class="docutils literal notranslate">axis_name=('data', 'model')</code>, where:</li> </ol> <ul class="simple"> <li><code class="docutils literal notranslate">'data'</code>: the mesh dimension used for data-parallel sharding of the batch dimension of inputs and activations.</li> <li><code class="docutils literal notranslate">'model'</code>: the mesh dimension used for sharding parameters of the model across devices.</li> </ul> <ol class="arabic simple" start="3"> <li>Make a simple utility function <code class="docutils literal notranslate">mesh_sharding</code> for generating a sharding object from the mesh and any layout.</li> </ol> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>from jax.sharding import Mesh, PartitionSpec, NamedSharding from jax.lax import with_sharding_constraint from jax.experimental import mesh_utils </pre></div> </div> </div> </div> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre># Create a mesh and annotate each axis with a name. device_mesh = mesh_utils.create_device_mesh((2, 4)) print(device_mesh) mesh = Mesh(devices=device_mesh, axis_names=('data', 'model')) print(mesh) def mesh_sharding(pspec: PartitionSpec) -> NamedSharding: return NamedSharding(mesh, pspec) </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre>[[CpuDevice(id=0) CpuDevice(id=1) CpuDevice(id=2) CpuDevice(id=3)] [CpuDevice(id=4) CpuDevice(id=5) CpuDevice(id=6) CpuDevice(id=7)]] Mesh('data': 2, 'model': 4) </pre></div> </div> </div> </div> </div> <div class="section" id="define-a-layer"> <h2>Define a layer<a class="headerlink" href="#define-a-layer" title="Permalink to this heading">#</a></h2> Before defining a simple model, create an example layer called <code class="docutils literal notranslate">DotReluDot</code> (by subclassing <code class="docutils literal notranslate">flax.linen.Module</code>). The layer creates two parameters <code class="docutils literal notranslate">W1</code> and <code class="docutils literal notranslate">W2</code> for dot product multiplication, and uses the <code class="docutils literal notranslate">jax.nn.relu</code> (ReLU) activation function in-between. To shard the parameters efficiently, apply the following APIs to annotate the parameters and intermediate variables: <ol class="arabic simple"> <li>Use <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_partitioning"><code class="docutils literal notranslate">flax.linen.with_partitioning</code></a> to decorate the initializer function when creating sub-layers or raw parameters.</li> <li>Apply <a class="reference external" href="https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html"><code class="docutils literal notranslate">jax.lax.with_sharding_constraint</code></a> (formerly, <code class="docutils literal notranslate">pjit.with_sharding_constraint</code>) to annotate intermediate variables like <code class="docutils literal notranslate">y</code> and <code class="docutils literal notranslate">z</code> to force a particular sharding pattern when the ideal constraint is known.</li> </ol> <ul class="simple"> <li>This step is optional, but can sometimes help auto-SPMD to partition efficiently. In the example below, the call is not required, because XLA will figure out the same sharding layout for <code class="docutils literal notranslate">y</code> and <code class="docutils literal notranslate">z</code> regardless.</li> </ul> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>class DotReluDot(nn.Module): depth: int dense_init: Callable = nn.initializers.xavier_normal() @nn.compact def __call__(self, x): y = nn.Dense(self.depth, kernel_init=nn.with_partitioning(self.dense_init, (None, 'model')), use_bias=False, # or overwrite with `bias_init` )(x) y = jax.nn.relu(y) # Force a local sharding annotation. y = with_sharding_constraint(y, mesh_sharding(PartitionSpec('data', 'model'))) W2 = self.param( 'W2', nn.with_partitioning(self.dense_init, ('model', None)), (self.depth, x.shape[-1])) z = jnp.dot(y, W2) # Force a local sharding annotation. z = with_sharding_constraint(z, mesh_sharding(PartitionSpec('data', None))) # Return a tuple to conform with the API `flax.linen.scan` as shown in the cell below. return z, None </pre></div> </div> </div> </div> Note that device axis names like <code class="docutils literal notranslate">'data'</code>, <code class="docutils literal notranslate">'model'</code> or <code class="docutils literal notranslate">None</code> are passed into both <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_partitioning"><code class="docutils literal notranslate">flax.linen.with_partitioning</code></a> and <a class="reference external" href="https://github.com/google/jax/blob/main/jax/_src/pjit.py#L1516"><code class="docutils literal notranslate">jax.lax.with_sharding_constraint</code></a> API calls. This refers to how each dimension of this data should be sharded — either across one of the device mesh dimensions, or not sharded at all. For example: <ul class="simple"> <li>When you define <code class="docutils literal notranslate">W1</code> with shape <code class="docutils literal notranslate">(x.shape[-1], self.depth)</code> and annotate as <code class="docutils literal notranslate">(None, 'model')</code>: <ul> <li>The first dimension (of length <code class="docutils literal notranslate">x.shape[-1]</code>) will be replicated across all devices.</li> <li>The second dimension (of length <code class="docutils literal notranslate">self.depth</code>) will be sharded over the <code class="docutils literal notranslate">'model'</code> axis of the device mesh. This means <code class="docutils literal notranslate">W1</code> will be sharded 4-way on devices <code class="docutils literal notranslate">(0, 4)</code>, <code class="docutils literal notranslate">(1, 5)</code>, <code class="docutils literal notranslate">(2, 6)</code> and <code class="docutils literal notranslate">(3, 7)</code>, on this dimension.</li> </ul> </li> <li>When you annotate the output <code class="docutils literal notranslate">z</code> as <code class="docutils literal notranslate">('data', None)</code>: <ul> <li>The first dimension — the batch dimension — will be sharded over the <code class="docutils literal notranslate">'data'</code> axis. This means half of the batch will be processed on devices <code class="docutils literal notranslate">0-3</code> (first four devices), and another half on devices <code class="docutils literal notranslate">4-7</code> (the remaining four devices).</li> <li>The second dimension — the data depth dimension — will be replicated across all devices.</li> </ul> </li> </ul> </div> <div class="section" id="define-a-model-with-flax-linen-scan-lifted-transformation"> <h2>Define a model with <code class="docutils literal notranslate">flax.linen.scan</code> lifted transformation<a class="headerlink" href="#define-a-model-with-flax-linen-scan-lifted-transformation" title="Permalink to this heading">#</a></h2> Having created <code class="docutils literal notranslate">DotReluDot</code>, you can now define the <code class="docutils literal notranslate">MLP</code> model (by subclassing <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module"><code class="docutils literal notranslate">flax.linen.Module</code></a>) as multiple layers of <code class="docutils literal notranslate">DotReluDot</code>. To replicate identical layers, you can either use <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/transformations.html#flax.linen.scan"><code class="docutils literal notranslate">flax.linen.scan</code></a>, or a for-loop: <ul class="simple"> <li><code class="docutils literal notranslate">flax.linen.scan</code> can provide faster compilation times.</li> <li>The for-loop can be faster on runtime.</li> </ul> The code below shows how to apply both methods, and default with the for-loop, so that all the parameters are two-dimensional and you can visualize their sharding. The <code class="docutils literal notranslate">flax.linen.scan</code> code is just to show that this API works with <a class="reference external" href="https://flax.readthedocs.io/en/latest/developer_notes/lift.html#supported-transformations">Flax lifted transforms</a>. <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>class MLP(nn.Module): num_layers: int depth: int use_scan: bool @nn.compact def __call__(self, x): if self.use_scan: x, _ = nn.scan(DotReluDot, length=self.num_layers, variable_axes={"params": 0}, split_rngs={"params": True}, metadata_params={nn.PARTITION_NAME: None} )(self.depth)(x) else: for i in range(self.num_layers): x, _ = DotReluDot(self.depth)(x) return x </pre></div> </div> </div> </div> Now, create a <code class="docutils literal notranslate">model</code> instance, and a sample input <code class="docutils literal notranslate">x</code>. <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre># MLP hyperparameters. BATCH, LAYERS, DEPTH, USE_SCAN = 8, 4, 1024, False # Create fake inputs. x = jnp.ones((BATCH, DEPTH)) # Initialize a PRNG key. k = random.key(0) # Create an Optax optimizer. optimizer = optax.adam(learning_rate=0.001) # Instantiate the model. model = MLP(LAYERS, DEPTH, USE_SCAN) </pre></div> </div> </div> </div> </div> <div class="section" id="specify-sharding"> <h2>Specify sharding<a class="headerlink" href="#specify-sharding" title="Permalink to this heading">#</a></h2> Next, you need to tell <code class="docutils literal notranslate">jax.jit</code> how to shard our data across devices. <div class="section" id="the-input-s-sharding"> <h3>The input’s sharding<a class="headerlink" href="#the-input-s-sharding" title="Permalink to this heading">#</a></h3> For data parallelism, you can shard the batched input <code class="docutils literal notranslate">x</code> across the <code class="docutils literal notranslate">data</code> axis by denoting the batch axis as <code class="docutils literal notranslate">'data'</code>. Then, use <a class="reference external" href="https://jax.readthedocs.io/en/latest/_autosummary/jax.device_put.html"><code class="docutils literal notranslate">jax.device_put</code></a> to place it onto the correct <code class="docutils literal notranslate">device</code>s. <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>x_sharding = mesh_sharding(PartitionSpec('data', None)) # dimensions: (batch, length) x = jax.device_put(x, x_sharding) jax.debug.visualize_array_sharding(x) </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output text_html"><pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> CPU 0,1,2,3 CPU 4,5,6,7 </pre> </div></div> </div> </div> <div class="section" id="the-output-s-sharding"> <h3>The output’s sharding<a class="headerlink" href="#the-output-s-sharding" title="Permalink to this heading">#</a></h3> You need to compile <code class="docutils literal notranslate">model.init()</code> (that is, <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init"><code class="docutils literal notranslate">flax.linen.Module.init()</code></a>), and its output as a pytree of parameters. Additionally, you may sometimes need wrap it with a <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.training.html#flax.training.train_state.TrainState"><code class="docutils literal notranslate">flax.training.train_state</code></a> to track other variables, such as optimizer states, and that would make the output an even more complex pytree. To achieve this, luckily, you don’t have to hardcode the output’s sharding by hand. Instead, you can: <ol class="arabic simple"> <li>Evaluate <code class="docutils literal notranslate">model.init</code> (in this case, a wrapper of it) abstractly using <a class="reference external" href="https://jax.readthedocs.io/en/latest/_autosummary/jax.eval_shape.html"><code class="docutils literal notranslate">jax.eval_shape</code></a>.</li> <li>Use <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.get_sharding"><code class="docutils literal notranslate">flax.linen.get_sharding</code></a> to automatically generate the <code class="docutils literal notranslate">jax.sharding.NamedSharding</code>. <ul class="simple"> <li>This step utilizes the <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_partitioning"><code class="docutils literal notranslate">flax.linen.with_partitioning</code></a> annotations in the earlier definition to generate the correct sharding for the parameters.</li> </ul> </li> </ol> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>def init_fn(k, x, model, optimizer): variables = model.init(k, x) # Initialize the model. state = train_state.TrainState.create( # Create a `TrainState`. apply_fn=model.apply, params=variables['params'], tx=optimizer) return state </pre></div> </div> </div> </div> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre># Create an abstract closure to wrap the function before feeding it in # because `jax.eval_shape` only takes pytrees as arguments. abstract_variables = jax.eval_shape( functools.partial(init_fn, model=model, optimizer=optimizer), k, x) # This `state_sharding` has the same pytree structure as `state`, the output # of the `init_fn`. state_sharding = nn.get_sharding(abstract_variables, mesh) state_sharding </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre>TrainState(step=NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(), memory_kind=unpinned_host), apply_fn=<bound method Module.apply of MLP( # attributes num_layers = 4 depth = 1024 use_scan = False )>, params={'DotReluDot_0': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None), memory_kind=unpinned_host)}, 'DotReluDot_1': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None), memory_kind=unpinned_host)}, 'DotReluDot_2': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None), memory_kind=unpinned_host)}, 'DotReluDot_3': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None), memory_kind=unpinned_host)}}, tx=GradientTransformationExtraArgs(init=<function chain.<locals>.init_fn at 0x7fc60ff228c0>, update=<function chain.<locals>.update_fn at 0x7fc60ff22290>), opt_state=(ScaleByAdamState(count=NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(), memory_kind=unpinned_host), mu={'DotReluDot_0': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None), memory_kind=unpinned_host)}, 'DotReluDot_1': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None), memory_kind=unpinned_host)}, 'DotReluDot_2': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None), memory_kind=unpinned_host)}, 'DotReluDot_3': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None), memory_kind=unpinned_host)}}, nu={'DotReluDot_0': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None), memory_kind=unpinned_host)}, 'DotReluDot_1': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None), memory_kind=unpinned_host)}, 'DotReluDot_2': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None), memory_kind=unpinned_host)}, 'DotReluDot_3': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None), memory_kind=unpinned_host)}}), EmptyState())) </pre></div> </div> </div> </div> </div> </div> <div class="section" id="compile-the-code"> <h2>Compile the code<a class="headerlink" href="#compile-the-code" title="Permalink to this heading">#</a></h2> Now you can apply <a class="reference external" href="https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html"><code class="docutils literal notranslate">jax.jit</code></a> to your <code class="docutils literal notranslate">init_fn</code>, but with two extra arguments: <code class="docutils literal notranslate">in_shardings</code> and <code class="docutils literal notranslate">out_shardings</code>. Run it to get the <code class="docutils literal notranslate">initialized_state</code>, in which parameters are sharded exactly as instructed: <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>jit_init_fn = jax.jit(init_fn, static_argnums=(2, 3), in_shardings=(mesh_sharding(()), x_sharding), # PRNG key and x out_shardings=state_sharding) initialized_state = jit_init_fn(k, x, model, optimizer) # for weight, partitioned in initialized_state.params['DotReluDot_0'].items(): # print(f'Sharding of {weight}: {partitioned.names}') jax.debug.visualize_array_sharding(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value) jax.debug.visualize_array_sharding(initialized_state.params['DotReluDot_0']['W2'].value) </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output text_html"><pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> CPU 0,4 CPU 1,5 CPU 2,6 CPU 3,7 </pre> </div><div class="output text_html"><pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> CPU 0,4 CPU 1,5 CPU 2,6 CPU 3,7 </pre> </div></div> </div> </div> <div class="section" id="inspect-the-module-output"> <h2>Inspect the Module output<a class="headerlink" href="#inspect-the-module-output" title="Permalink to this heading">#</a></h2> Note that in the output of <code class="docutils literal notranslate">initialized_state</code>, the <code class="docutils literal notranslate">params</code> <code class="docutils literal notranslate">W1</code> and <code class="docutils literal notranslate">W2</code> are of type <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.Partitioned"><code class="docutils literal notranslate">flax.linen.Partitioned</code></a>. This is a wrapper around the actual <code class="docutils literal notranslate">jax.Array</code> that allows Flax to record the axis names associated with it. You can access the raw <code class="docutils literal notranslate">jax.Array</code>s by calling <code class="docutils literal notranslate">flax.linen.meta.unbox()</code> upon the dictionary, or call <code class="docutils literal notranslate">.value</code> upon individual variable. You can also use <code class="docutils literal notranslate">flax.linen.meta.replace_boxed()</code> to change the underlying <code class="docutils literal notranslate">jax.Array</code> without modifying the sharding annotations. <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>print(type(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'])) print(type(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value)) print(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].names) print(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value.shape) </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><class 'flax.core.meta.Partitioned'> <class 'jaxlib.xla_extension.ArrayImpl'> (None, 'model') (1024, 1024) </pre></div> </div> </div> </div> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre># Say for some unknown reason you want to make the whole param tree all-zero unboxed_params = nn.meta.unbox(initialized_state.params) all_zero = jax.tree.map(jnp.zeros_like, unboxed_params) all_zero_params = nn.meta.replace_boxed(initialized_state.params, all_zero) assert jnp.sum(nn.meta.unbox(all_zero_params['DotReluDot_0']['Dense_0']['kernel'])) == 0 </pre></div> </div> </div> </div> You can also check the underlying <a class="reference external" href="https://jax.readthedocs.io/en/latest/jax.sharding.html"><code class="docutils literal notranslate">jax.sharding</code></a> of each parameter, which is now more internal than <code class="docutils literal notranslate">NamedSharding</code>. Note that numbers like <code class="docutils literal notranslate">initialized_state.step</code> are replicated across all devices. <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value.sharding </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre>NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host) </pre></div> </div> </div> </div> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>print(initialized_state.step) initialized_state.step.sharding </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre>0 </pre></div> </div> <div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre>NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(), memory_kind=unpinned_host) </pre></div> </div> </div> </div> You can use <a class="reference external" href="https://jax.readthedocs.io/en/latest/_autosummary/jax.tree_util.tree_map.html"><code class="docutils literal notranslate">jax.tree_util.tree_map</code></a> to perform mass computation on a dict of boxed params, in the same way as on a dict of JAX arrays. <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>diff = jax.tree_util.tree_map( lambda a, b: a - b, initialized_state.params['DotReluDot_0'], initialized_state.params['DotReluDot_0']) print(jax.tree_util.tree_map(jnp.shape, diff)) diff_array = diff['Dense_0']['kernel'].value print(type(diff_array)) print(diff_array.shape) </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre>{'Dense_0': {'kernel': Partitioned(value=(1024, 1024), names=(None, 'model'), mesh=None)}, 'W2': Partitioned(value=(1024, 1024), names=('model', None), mesh=None)} <class 'jaxlib.xla_extension.ArrayImpl'> (1024, 1024) </pre></div> </div> </div> </div> </div> <div class="section" id="compile-the-train-step-and-inference"> <h2>Compile the train step and inference<a class="headerlink" href="#compile-the-train-step-and-inference" title="Permalink to this heading">#</a></h2> Create a <code class="docutils literal notranslate">jit</code>ted training step as follows: <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>@functools.partial(jax.jit, in_shardings=(state_sharding, x_sharding), out_shardings=state_sharding) def train_step(state, x): # A fake loss function. def loss_unrolled(params): y = model.apply({'params': params}, x) return y.sum() grad_fn = jax.grad(loss_unrolled) grads = grad_fn(state.params) state = state.apply_gradients(grads=grads) return state with mesh: new_state = train_step(initialized_state, x) </pre></div> </div> </div> </div> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>print(f'Sharding of Weight 1:') jax.debug.visualize_array_sharding(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value) print(f'Sharding of Weight 2:') jax.debug.visualize_array_sharding(initialized_state.params['DotReluDot_0']['W2'].value) </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre>Sharding of Weight 1: </pre></div> </div> <div class="output text_html"><pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> CPU 0,4 CPU 1,5 CPU 2,6 CPU 3,7 </pre> </div><div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre>Sharding of Weight 2: </pre></div> </div> <div class="output text_html"><pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> CPU 0,4 CPU 1,5 CPU 2,6 CPU 3,7 </pre> </div></div> </div> Then, create a compiled inference step. Note that the output is also sharded along <code class="docutils literal notranslate">(data, None)</code>. <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>@functools.partial(jax.jit, in_shardings=(state_sharding, x_sharding), out_shardings=x_sharding) def apply_fn(state, x): return state.apply_fn({'params': state.params}, x) with mesh: y = apply_fn(new_state, x) print(type(y)) print(y.dtype) print(y.shape) jax.debug.visualize_array_sharding(y) </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre><class 'jaxlib.xla_extension.ArrayImpl'> float32 (8, 1024) </pre></div> </div> <div class="output text_html"><pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> CPU 0,1,2,3 CPU 4,5,6,7 </pre> </div></div> </div> </div> <div class="section" id="profiling"> <h2>Profiling<a class="headerlink" href="#profiling" title="Permalink to this heading">#</a></h2> If you are running on a TPU pod or a pod slice, you can use a custom <code class="docutils literal notranslate">block_all</code> utility function, as defined below, to measure the performance: <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>%%timeit def block_all(xs): jax.tree_util.tree_map(lambda x: x.block_until_ready(), xs) return xs with mesh: new_state = block_all(train_step(initialized_state, x)) </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre>237 ms ± 9.08 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) </pre></div> </div> </div> </div> </div> <div class="section" id="logical-axis-annotation"> <h2>Logical axis annotation<a class="headerlink" href="#logical-axis-annotation" title="Permalink to this heading">#</a></h2> JAX’s automatic SPMD encourages users to explore different sharding layouts to find the optimal one. To this end, in Flax you actually can annotate the dimensions of any data with more descriptive axis names (not just device mesh axis names like <code class="docutils literal notranslate">'data'</code> and <code class="docutils literal notranslate">'model'</code>). The <code class="docutils literal notranslate">LogicalDotReluDot</code> and <code class="docutils literal notranslate">LogicalMLP</code> Module definition below are similar to the Modules you created earlier, except for the following: <ol class="arabic simple"> <li>All axes are annotated with more concrete, meaningful names, such as <code class="docutils literal notranslate">'embed'</code>, <code class="docutils literal notranslate">'hidden'</code>, <code class="docutils literal notranslate">'batch'</code> and <code class="docutils literal notranslate">'layer'</code>. These names are referred to as logical axis names in Flax. They make the dimensional changes inside model definitions more readable.</li> <li><a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_logical_partitioning"><code class="docutils literal notranslate">flax.linen.with_logical_partitioning</code></a> replaces <code class="docutils literal notranslate">flax.linen.with_partitioning</code>; and <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_logical_constraint"><code class="docutils literal notranslate">flax.linen.with_logical_constraint</code></a> replaces <code class="docutils literal notranslate">jax.lax.with_sharding_constraint</code>, to recognize the logical axis names.</li> </ol> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>class LogicalDotReluDot(nn.Module): depth: int dense_init: Callable = nn.initializers.xavier_normal() @nn.compact def __call__(self, x): y = nn.Dense(self.depth, kernel_init=nn.with_logical_partitioning(self.dense_init, ('embed', 'hidden')), use_bias=False, # or overwrite with `bias_init` )(x) y = jax.nn.relu(y) # Force a local sharding annotation. y = with_sharding_constraint(y, mesh_sharding(PartitionSpec('data', 'model'))) W2 = self.param( 'W2', nn.with_logical_partitioning(self.dense_init, ('hidden', 'embed')), (self.depth, x.shape[-1])) z = jnp.dot(y, W2) # Force a local sharding annotation. z = nn.with_logical_constraint(z, ('batch', 'embed')) return z, None class LogicalMLP(nn.Module): num_layers: int depth: int use_scan: bool @nn.compact def __call__(self, x): if self.use_scan: x, _ = nn.scan(LogicalDotReluDot, length=self.num_layers, variable_axes={"params": 0}, split_rngs={"params": True}, metadata_params={nn.PARTITION_NAME: 'layer'} )(self.depth)(x) else: for i in range(self.num_layers): x, _ = LogicalDotReluDot(self.depth)(x) return x </pre></div> </div> </div> </div> Now, initiate a model and try to figure out what sharding its <code class="docutils literal notranslate">state</code> should have. To allow the device mesh to take your model correctly, you need to decide which of these logical axis names are mapped to the device axis <code class="docutils literal notranslate">'data'</code> or <code class="docutils literal notranslate">'model'</code>. This rule is a list of (<code class="docutils literal notranslate">logical_axis_name</code>, <code class="docutils literal notranslate">device_axis_name</code>) tuples, and <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.logical_to_mesh_sharding"><code class="docutils literal notranslate">flax.linen.logical_to_mesh_sharding</code></a> will convert them to the kind of sharding that the device mesh can understand. This allows you to change the rules and try out new partition layouts without modifying the model definition. <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre># Unspecified rule means unsharded by default, so no need to specify `('embed', None)` and `('layer', None)`. rules = (('batch', 'data'), ('hidden', 'model')) logical_model = LogicalMLP(LAYERS, DEPTH, USE_SCAN) logical_abstract_variables = jax.eval_shape( functools.partial(init_fn, model=logical_model, optimizer=optimizer), k, x) logical_state_spec = nn.get_partition_spec(logical_abstract_variables) print('annotations are logical, not mesh-specific: ', logical_state_spec.params['LogicalDotReluDot_0']['Dense_0']['kernel']) logical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, mesh, rules) print('sharding annotations are mesh-specific: ', logical_state_sharding.params['LogicalDotReluDot_0']['Dense_0']['kernel'].spec) </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre>annotations are logical, not mesh-specific: PartitionSpec('embed', 'hidden') sharding annotations are mesh-specific: PartitionSpec(None, 'model') </pre></div> </div> </div> </div> You can verify that the <code class="docutils literal notranslate">logical_state_spec</code> here has the same content as <code class="docutils literal notranslate">state_spec</code> in the previous (“non-logical”) example. This allows you to <code class="docutils literal notranslate">jax.jit</code> your Module’s <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init"><code class="docutils literal notranslate">flax.linen.Module.init</code></a> and <a class="reference external" href="https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.apply"><code class="docutils literal notranslate">flax.linen.Module.apply</code></a> the same way in the above above. <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>state_sharding.params['DotReluDot_0'] == logical_state_sharding.params['LogicalDotReluDot_0'] </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre>True </pre></div> </div> </div> </div> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>logical_jit_init_fn = jax.jit(init_fn, static_argnums=(2, 3), in_shardings=(mesh_sharding(()), x_sharding), # PRNG key and x out_shardings=logical_state_sharding) logical_initialized_state = logical_jit_init_fn(k, x, logical_model, optimizer) </pre></div> </div> </div> </div> <div class="cell docutils container"> <div class="cell_input docutils container"> <div class="highlight-ipython3 notranslate"><div class="highlight"><pre>print(f'Sharding of Weight 1:') jax.debug.visualize_array_sharding(logical_initialized_state.params['LogicalDotReluDot_0']['Dense_0']['kernel'].value) print(f'Sharding of Weight 2:') jax.debug.visualize_array_sharding(logical_initialized_state.params['LogicalDotReluDot_0']['W2'].value) </pre></div> </div> </div> <div class="cell_output docutils container"> <div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre>Sharding of Weight 1: </pre></div> </div> <div class="output text_html"><pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> CPU 0,4 CPU 1,5 CPU 2,6 CPU 3,7 </pre> </div><div class="output stream highlight-myst-ansi notranslate"><div class="highlight"><pre>Sharding of Weight 2: </pre></div> </div> <div class="output text_html"><pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> CPU 0,4 CPU 1,5 CPU 2,6 CPU 3,7 </pre> </div></div> </div> </div> <div class="section" id="when-to-use-device-axis-logical-axis"> <h2>When to use device axis / logical axis<a class="headerlink" href="#when-to-use-device-axis-logical-axis" title="Permalink to this heading">#</a></h2> Choosing when to use a device or logical axis depends on how much you want to control the partitioning of your model: <ul class="simple"> <li>Device mesh axis: If you want a very simple model, or you are very confident of your way of partitioning, defining it with device mesh axis can potentially save you a few extra lines of code of converting the logical naming back to the device naming.</li> <li>Logical naming: On the other hand, the logical naming helpers can be useful for exploring different sharding layouts. Use this if you want to experiment around and find the most optimal partition layout for your model.</li> <li>Device axis names: In really advanced use cases, you may have more complicated sharding patterns that require annotating activation dimension names differently from parameter dimension names. If you wish to have more fine-grained control on manual mesh assignments, directly using device axis names could be more helpful.</li> </ul> </div> <div class="section" id="save-the-data"> <h2>Save the data<a class="headerlink" href="#save-the-data" title="Permalink to this heading">#</a></h2> To save the cross-device array, you can use <a class="reference external" href="https://flax.readthedocs.io/en/latest/_modules/flax/training/checkpoints.html"><code class="docutils literal notranslate">flax.training.checkpoints</code></a>, as shown in the <a class="reference external" href="https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html#multi-host-multi-process-checkpointing">Save and load checkpoints guide - Multi-host/multi-process checkpointing</a>. This is especially required if you are running on a multi-host environment (for example, a TPU pod). In practice, you might want to save the raw <code class="docutils literal notranslate">jax.Array</code> pytree as checkpoint, instead of the wrapped <code class="docutils literal notranslate">Partitioned</code> values, to reduce complexity. You can restore it as-is and put it back into an annotated pytree with <code class="docutils literal notranslate">flax.linen.meta.replace_boxed()</code>. Keep in mind that to restore the arrays to the desired partition, you need to provide a sample <code class="docutils literal notranslate">target</code> pytree that has the same structure and has the desired <a class="reference external" href="https://jax.readthedocs.io/en/latest/jax.sharding.html#jax.sharding.Sharding"><code class="docutils literal notranslate">jax.sharding.Sharding</code></a> in place for each JAX array. The sharding you use to restore the array doesn’t necessarily need to be the same as the ones you used to store the array. </div> </div> </article> <footer class="prev-next-footer d-print-none"> <div class="prev-next-area"> <a class="left-prev" href="ensembling.html" title="previous page"> <div class="prev-next-info"> previous Ensembling on multiple devices </div> </a> <a class="right-next" href="../model_inspection/index.html" title="next page"> <div class="prev-next-info"> next Model inspection </div> </a> </div> </footer> </div> <div class="bd-sidebar-secondary bd-toc"><div class="sidebar-secondary-items sidebar-secondary__inner"> <div class="sidebar-secondary-item"> <div class="page-toc tocsection onthispage"> Contents </div> <nav class="bd-toc-nav page-toc"> <ul class="visible nav section-nav flex-column"> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#flax-and-jax-jit-scaled-up">Flax and <code class="docutils literal notranslate">jax.jit</code> scaled up</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#setup">Setup</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#define-a-layer">Define a layer</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#define-a-model-with-flax-linen-scan-lifted-transformation">Define a model with <code class="docutils literal notranslate">flax.linen.scan</code> lifted transformation</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#specify-sharding">Specify sharding</a><ul class="nav section-nav flex-column"> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#the-input-s-sharding">The input’s sharding</a></li> <li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#the-output-s-sharding">The output’s sharding</a></li> </ul> </li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#compile-the-code">Compile the code</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#inspect-the-module-output">Inspect the Module output</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#compile-the-train-step-and-inference">Compile the train step and inference</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#profiling">Profiling</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#logical-axis-annotation">Logical axis annotation</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#when-to-use-device-axis-logical-axis">When to use device axis / logical axis</a></li> <li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#save-the-data">Save the data</a></li> </ul> </nav></div> </div></div> </div> <footer class="bd-footer-content"> <div class="bd-footer-content__inner container"> <div class="footer-item"> By The Flax authors </div> <div class="footer-item"> © Copyright 2023, The Flax authors. </div> <div class="footer-item"> </div> <div class="footer-item"> </div> </div> </footer> </main> </div> </div>  <script src="../../_static/scripts/bootstrap.js?digest=dfe6caa3a7d634c4db9b"></script> <script src="../../_static/scripts/pydata-sphinx-theme.js?digest=dfe6caa3a7d634c4db9b"></script> <footer class="bd-footer"> </footer> </body> </html>