Skip to content

[NNX] Delete Linen (12.1a): collapse dispatch in the core train loop and maxtext_utils #4683

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

Draft
ecnal-cienet wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[NNX] Delete Linen (12.1a): collapse dispatch in the core train loop and maxtext_utils #4683

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
63 changes: 5 additions & 58 deletions src/maxtext/utils/maxtext_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,10 +96,7 @@ def get_functional_train_with_signature(
"""Get the shardings (both state and data) for `train_step`."""
functional_train = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings)
functional_train.__name__ = "train_step" # pyrefly: ignore[missing-attribute]
if config.pure_nnx:
in_shardings = (state_mesh_shardings, data_sharding) # State, batch
else:
in_shardings = (state_mesh_shardings, data_sharding, None) # State, batch, rng
in_shardings = (state_mesh_shardings, data_sharding) # State, batch
out_shardings = (state_mesh_shardings, None) # State, metrics
static_argnums = () # We partial out the static argnums of model and config
donate_argnums = 0 # This is the index of the state - we allow the compiler to make use of this memory.
Expand All @@ -110,10 +107,7 @@ def get_functional_eval_with_signature(eval_step, data_sharding, state_mesh_shar
"""Get the shardings (both state and data) for `eval_step`."""
functional_eval = functools.partial(eval_step, model, config)
functional_eval.__name__ = "eval_step" # pyrefly: ignore[missing-attribute]
if config.pure_nnx:
in_shardings = (state_mesh_shardings, data_sharding) # State, batch (NNX: no rng)
else:
in_shardings = (state_mesh_shardings, data_sharding, None) # State, batch, rng
in_shardings = (state_mesh_shardings, data_sharding) # State, batch
out_shardings = None # metrics
static_argnums = () # We partial out the static argnums of model, config
donate_argnums = () # state will be kept instead of being donated in eval_step
Expand Down Expand Up @@ -254,11 +248,7 @@ def get_train_input_output_trees(func, input_args, input_kwargs):

serialized_compiled = load_serialized_compiled(config.compiled_trainstep_file)
shaped_batch = get_shaped_batch(config)
if config.pure_nnx:
shaped_input_args = (state, shaped_batch)
else:
example_rng = jax.random.PRNGKey(0)
shaped_input_args = (state, shaped_batch, example_rng)
shaped_input_args = (state, shaped_batch)
shaped_input_kwargs = {}
in_tree, out_tree = get_train_input_output_trees(partial_train, shaped_input_args, shaped_input_kwargs)
p_train_step = deserialize_and_load(serialized_compiled, in_tree, out_tree, execution_devices=execution_devices)
Expand Down Expand Up @@ -1794,51 +1784,8 @@ def get_logical_annotations(config, mesh, init_state_fn):


def get_abstract_state(config, mesh, init_state_fn, is_training=True):
"""Get a shaped abstraction of the state (including optimizer)"""
if config.pure_nnx:
return get_abstract_state_nnx(config, mesh, init_state_fn, is_training)

init_state_partial = init_state_fn

with nn_partitioning.axis_rules(config.logical_axis_rules):
abstract_state = jax.eval_shape(init_state_partial)

state_logical_annotations = nn.get_partition_spec(abstract_state)

state_mesh_shardings = nn.logical_to_mesh_sharding(state_logical_annotations, mesh, config.logical_axis_rules)
if is_training and config.shard_optimizer_over_data:
# Add data to sharding for optimizer state
state_mesh_shardings = state_mesh_shardings.replace(
opt_state=jax.tree.map_with_path(
functools.partial(sharding.add_data_to_sharding, mesh),
max_utils.unbox_logicallypartioned(abstract_state).opt_state,
state_mesh_shardings.opt_state,
)
)
if is_training and config.optimizer_memory_host_offload:
opt_state = jax.tree_util.tree_map(lambda x: x.with_memory_kind(kind="pinned_host"), state_mesh_shardings.opt_state)
state_mesh_shardings = state_mesh_shardings.replace(opt_state=opt_state)
if is_training and config.parameter_memory_host_offload:
assert config.param_scan_axis == 0, "You must set the scan axis 0 to enable parameter offloading."

def move(path, x):
max_logging.log(f"max_utils.py: Moving {path} to host")
return x.with_memory_kind(kind="pinned_host")

params = jax.tree_util.tree_map_with_path(move, state_mesh_shardings.params)
state_mesh_shardings = state_mesh_shardings.replace(params=params)

abstract_sharded_state = jax.jit(init_state_partial, in_shardings=None, out_shardings=state_mesh_shardings).eval_shape()

unboxed_abstract_sharded_state = max_utils.unbox_logicallypartioned(abstract_sharded_state)
# Initialization
with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):
state_mesh_annotations = nn.logical_to_mesh(state_logical_annotations)
return (
unboxed_abstract_sharded_state,
state_mesh_annotations,
state_mesh_shardings,
)
"""Get a shaped abstraction of the state (including optimizer)."""
return get_abstract_state_nnx(config, mesh, init_state_fn, is_training)


def get_abstract_state_nnx(config, mesh, nnx_init_trainstate_fn, is_training=True):
Expand Down
83 changes: 30 additions & 53 deletions src/maxtext/utils/train_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
import jax
import functools
import orbax.checkpoint.pathways as ocp_pathways
from functools import partial

from flax import nnx
from flax.linen import partitioning as nn_partitioning
Expand Down Expand Up @@ -228,27 +227,20 @@ def setup_train_loop(config, recorder, devices=None):
from maxtext.input_pipeline.input_pipeline_interface import create_data_iterator

with maybe_record_goodput(recorder, GoodputEvent.TPU_INIT):
is_training = True
init_rng = jax.random.PRNGKey(config.init_weights_seed)
mesh = maxtext_utils.get_mesh_from_config(config, devices)
context_parallel_size = mesh.shape.get(config.context_sharding, 1)
if config.pure_nnx:
# Create abstract NNX model.
_create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices)
else:
model = model_creation_utils.from_config(config, devices)
# Create abstract NNX model.
_create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices)
learning_rate_schedule, tx = create_training_optimizer(config, model)

if config.pure_nnx:
# For NNX, the train state is wrapped in the TrainStateNNX module.
def create_train_state_fn():
model = _create_model_partial()
optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)
return train_state_nnx.TrainStateNNX(model, optimizer)
# The train state is wrapped in the TrainStateNNX module.
def create_train_state_fn():
model = _create_model_partial()
optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)
return train_state_nnx.TrainStateNNX(model, optimizer)

init_state_fn = create_train_state_fn
else:
init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, is_training, init_rng)
init_state_fn = create_train_state_fn
checkpoint_manager = create_checkpoint_manager(config, mesh, init_state_fn)
if checkpoint_manager is not None:
checkpoint_step = checkpointing.latest_step(checkpoint_manager)
Expand Down Expand Up @@ -308,38 +300,30 @@ def create_train_state_fn():
state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state(
data_iterator, config, mesh, checkpoint_manager, init_state_fn
)
if config.pure_nnx:
with nn_partitioning.axis_rules(config.logical_axis_rules):
# We only need the graphdef here; it's merged with state below. Avoid
# nnx.get_abstract_model: it eagerly builds a NamedSharding for every variable
# under jax.set_mesh(mesh) and rejects any logical name missing from
# logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes
# without a mesh skips sharding resolution, so it avoids the crash.
state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn))
_, state_params, _ = nnx.split(state.model, nnx.Param, ...)
_, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...)
else:
state_params = state.params
state_mesh_shardings_params = state_mesh_shardings.params
with nn_partitioning.axis_rules(config.logical_axis_rules):
# We only need the graphdef here; it's merged with state below. Avoid
# nnx.get_abstract_model: it eagerly builds a NamedSharding for every variable
# under jax.set_mesh(mesh) and rejects any logical name missing from
# logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes
# without a mesh skips sharding resolution, so it avoids the crash.
state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn))
_, state_params, _ = nnx.split(state.model, nnx.Param, ...)
_, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...)

if config.enable_diloco:
with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):
state, outer_opt_state_sharding = diloco.build_diloco_state(config, lambda: state, mesh=mesh)

# create state_mesh_shardings for the DilocoState
step_mesh = state_mesh_shardings.optimizer.step.mesh if config.pure_nnx else state_mesh_shardings.step.mesh
step_mesh = state_mesh_shardings.optimizer.step.mesh
inner_state_shardings = diloco.add_diloco_to_sharding(state_mesh_shardings)
state_mesh_shardings = diloco.DiLoCoTrainState(
inner_state_shardings,
# Match the outer params' pure-dict structure (build_diloco_state stores
# outer_params via to_pure_dict), so the sharding tree matches the state tree.
state_mesh_shardings_params.to_pure_dict() # pyrefly: ignore[missing-attribute]
if config.pure_nnx
else state_mesh_shardings_params,
state_mesh_shardings_params.to_pure_dict(),
outer_opt_state_sharding,
jax.sharding.NamedSharding( # pyrefly: ignore[bad-argument-type]
mesh=step_mesh, spec=jax.sharding.PartitionSpec()
),
jax.sharding.NamedSharding(mesh=step_mesh, spec=jax.sharding.PartitionSpec()),
)

# TODO(aireenmei, hengtaoguo): support sharding in vit for multimodal
Expand All @@ -349,28 +333,21 @@ def create_train_state_fn():

# print weights sharding info under debug sharding mode
if config.debug_sharding:
if config.pure_nnx:
# TODO: Study how to get logical annotations of NNX module. Because of eager sharding, we
# probably already lost the logical partition info at this moment.
logical_annotations_params = None
else:
logical_annotations = maxtext_utils.get_logical_annotations(config, mesh, init_state_fn)
logical_annotations_params = logical_annotations.params
# TODO: Study how to get logical annotations of NNX module. Because of eager sharding, we
# probably already lost the logical partition info at this moment.
logical_annotations_params = None

max_utils.print_non_trivial_mesh_axis(model.mesh) # pyrefly: ignore[missing-attribute]
maxtext_utils.print_shardings_params(state_params, state_mesh_shardings_params, mesh, logical_annotations_params)

if config.pure_nnx:
if config.enable_diloco:
# Don't merge the DiLoCoTrainState into the plain-model graphdef. The inner
# train step needs that graphdef as jit_model; the wrapper passes through as state.
train_state = state
model = state_graphdef # pyrefly: ignore[unbound-name]
else:
train_state = nnx.merge(state_graphdef, state) # pyrefly: ignore[unbound-name]
model = train_state.model
else:
if config.enable_diloco:
# Don't merge the DiLoCoTrainState into the plain-model graphdef. The inner
# train step needs that graphdef as jit_model; the wrapper passes through as state.
train_state = state
model = state_graphdef
else:
train_state = nnx.merge(state_graphdef, state)
model = train_state.model

return (
init_rng,
Expand Down
7 changes: 3 additions & 4 deletions tests/integration/setup_train_loop_nnx_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

"""Integration test for setup_train_loop with pure_nnx=True.
"""Integration test for setup_train_loop on the NNX path.

setup_train_loop wires together create_nnx_abstract_model, the training
optimizer,
Expand Down Expand Up @@ -43,7 +43,6 @@ def _tiny_nnx_pyconfig(**overrides):
"enable_checkpointing": False,
"dataset_type": "synthetic",
"model_name": "default",
"pure_nnx": True,
"per_device_batch_size": 1.0,
"base_emb_dim": 8,
"base_num_query_heads": 4,
Expand All @@ -68,7 +67,7 @@ def _tiny_nnx_pyconfig(**overrides):
class SetupTrainLoopNNXIntegrationTest(unittest.TestCase):
"""End-to-end check that setup_train_loop returns a usable TrainStateNNX."""

def test_pure_nnx_setup_returns_train_state_nnx(self):
def test_setup_returns_train_state_nnx(self):
config = _tiny_nnx_pyconfig()

(
Expand Down Expand Up @@ -109,7 +108,7 @@ def test_pure_nnx_setup_returns_train_state_nnx(self):
# flag them as unused — they're part of the public return contract.
del checkpoint_manager, rampup_manager, eval_data_iterator

def test_pure_nnx_setup_param_only_split_matches_model(self):
def test_setup_param_only_split_matches_model(self):
"""nnx.split(state.model, nnx.Param, ...) must yield a non-empty Param tree

whose structure matches state_mesh_shardings.model after the same split.
Expand Down
Loading
Loading