Skip to content

[Jax API update] Remove jax_spmd_mode #1136

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 22 commits into
base: main
Choose a base branch
from
Open

[Jax API update] Remove jax_spmd_mode #1136

Show file tree
Hide file tree
Changes from 20 commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
835f0bd
modify primitives
Steboss Apr 1, 2025
dd380a4
fix pre-commit error
Steboss Apr 1, 2025
5f1640d
update jax.tree_map to jax.tree_util.tree_map
Steboss Apr 3, 2025
deb1da9
Update axlearn/common/array_serialization.py
Steboss Apr 3, 2025
4784822
fix black formatting
Steboss Apr 7, 2025
e1828ee
fix tree_map
Steboss Apr 11, 2025
87c23a0
update main branch
Steboss Apr 29, 2025
9d83d7c
conflicts to fix
Steboss Apr 29, 2025
a089d43
rebase with main to solve conflicts
Steboss Apr 11, 2025
13a8451
modify R0917
Steboss Apr 11, 2025
8bc6f64
modify R0917
Steboss Apr 11, 2025
16e4c04
Fix comments
Steboss Apr 15, 2025
9316c5c
use flatten_one_level_with_keys
Steboss Apr 16, 2025
8de926a
the spmd has been removed in commit 7634230cdcd2d3cb42d1093f6ab255f47…
Steboss Apr 28, 2025
c3a46d4
fix the array_serialization
Steboss Apr 29, 2025
0ddb457
fix conflicts for optimizers.py
Steboss Apr 29, 2025
835ab56
Update axlearn/common/utils.py
Steboss Apr 29, 2025
2c56be8
check pylint
Steboss Apr 29, 2025
d2774c3
fix docstring
Steboss Apr 29, 2025
7e9ca87
use public API
Steboss Apr 30, 2025
c860d7f
fix tryexcept
Steboss May 1, 2025
0605432
make sure pytree contains all the previous cases + fix tests
Steboss May 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 6 additions & 19 deletions axlearn/common/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,6 @@

import attr
import jax
import jax.flatten_util
import numpy as np
from absl import logging
from jax import numpy as jnp
Expand All @@ -54,6 +53,7 @@
from jax.experimental import mesh_utils, multihost_utils
from jax.extend.core import Primitive
from jax.sharding import PartitionSpec
from jax.tree_util import default_registry

from axlearn.common import serialization
from axlearn.common.config import (
Expand Down Expand Up @@ -1853,36 +1853,23 @@ def thread_stack_traces() -> Sequence[Sequence[str]]:
return grouped_lines


def pytree_children(node: Any) -> Sequence[tuple[KeyEntry, Any]]:
def pytree_children(node: Any) -> list[tuple[KeyEntry, Any]]:
Copy link
Contributor

@apghml apghml May 2, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this example from the jax docs fail with the new implementation?

import jax.numpy as jnp
import jax.tree
from jax.tree_util import GetAttrKey


class MyContainer:

    def __init__(self):

        self.x = jnp.zeros(2)
        self.y = jnp.ones(2)



def flatten_with_keys(obj):

    children = [(GetAttrKey('x'), obj.x),
                (GetAttrKey('y'), obj.y)]



    aux_data = ()  # aux_data must contain static, hashable data.

    return children, aux_data

def unflatten(aux_data, children):
    obj = object.__new__(MyContainer)

    obj.x, obj.y = children

    obj.size, = aux_data

    return obj

jax.tree_util.register_pytree_node(MyContainer, flatten_with_keys, unflatten)

pytree_children(MyContainer())

"""Generate the (key, value) pairs for the immediate children of a pytree `node`.

The returned children match those returned by
`jax.tree_util.default_registry.flatten_one_level()`.

Reference: jax._src.tree_util.generate_key_paths()

Example:
```
assert pytree_children(dict(a=[1,2])) == [(DictKey('a'), [1,2])]
```
"""
# pylint: disable-next=protected-access
registry_with_keypaths = jax._src.tree_util._registry_with_keypaths

key_handler = registry_with_keypaths.get(type(node))
if key_handler:
key_children, _ = key_handler.flatten_with_keys(node)
return key_children

flat = jax.tree_util.default_registry.flatten_one_level(node)
if flat is None:
try:
key_child_pairs, _ = default_registry.flatten_one_level_with_keys(node)
return list(key_child_pairs)
except ValueError:
return []

if isinstance(node, tuple) and hasattr(node, "_fields") and flat[1] == type(node):
# Handle namedtuple as a special case, based on heuristic.
return [(jax.tree_util.GetAttrKey(s), getattr(node, s)) for s in node._fields]
return [(jax.tree_util.FlattenedIndexKey(i), c) for i, c in enumerate(flat[0])]


def find_cycles(tree: Nested) -> dict[str, KeyPath]:
"""Find a cycle in pytree `tree` if one exists.
Expand Down
2 changes: 0 additions & 2 deletions axlearn/common/utils_spmd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,6 @@ def setup(
"""
# Use a GSPMD-friendly PRNG implementation.
jax.config.update("jax_default_prng_impl", "rbg")
# This allows replicated jax.Arrays to be used for computation on the host.
jax.config.update("jax_spmd_mode", "allow_all")

global _jax_distributed_initialized # pylint: disable=global-statement
if not _jax_distributed_initialized:
Expand Down