Merge pull request #2511 from AI-Hypercomputer:fix_linter

PiperOrigin-RevId: 820348631

Author: Google-ML-Automation

src/MaxText/maxtext_utils.py

@@ -53,7 +53,9 @@ def get_input_data_sharding(config, mesh):
   return nn.logical_to_mesh_sharding(P(*config.input_data_sharding_logical_axes), mesh, config.logical_axis_rules)
 
 
-def get_functional_train_with_signature(train_step, data_sharding, state_mesh_shardings, model, config, params_shardings=None):
+def get_functional_train_with_signature(
+    train_step, data_sharding, state_mesh_shardings, model, config, params_shardings=None
+):
   """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"
@@ -97,14 +99,14 @@ def get_shaped_batch(config):
 
 def should_prevent_cse_in_remat(config):
   """Determines whether to prevent common subexpression elimination (CSE) in remat.
-  
+
   CSE should not be prevented when:
   1. Layers are being scanned (scan_layers=True), OR
   2. Gradient accumulation is enabled (gradient_accumulation_steps > 1) on GPU hardware
-  
+
   Args:
     config: Configuration object with scan_layers, gradient_accumulation_steps, and hardware
-  
+
   Returns:
     bool: True if CSE should be prevented, False otherwise
   """
@@ -961,6 +963,7 @@ def setup_training_state(model, data_iterator, tx, config, rng, mesh, checkpoint
       is_training,
   )
 
+
 def unbox_logicallypartioned(boxed_pytree):
   """Unboxes the flax.LogicallyPartitioned pieces
   Args:
@@ -975,23 +978,24 @@ def unbox_logicallypartioned(boxed_pytree):
       is_leaf=lambda k: isinstance(k, nn.spmd.LogicallyPartitioned),
   )
 
+
 def add_data_to_sharding(mesh, path, aval, sharding):
   """Adds 'data' dimension to sharding spec if compatible and not already present.
-  
+
   This function attempts to add data parallelism to a sharding specification by finding
   a dimension that is divisible by the 'data' mesh axis size and doesn't conflict with
   existing partitioning (e.g., tensor parallelism).
   This function is mainly used to add data parallelism to the optimizer state for Zero-1 style sharding.
-  
+
   Args:
     mesh: The device mesh
     path: JAX tree path to the value being sharded
     aval: Abstract value with shape information
     sharding: Current NamedSharding to potentially augment
-  
+
   Returns:
     NamedSharding: Updated sharding with 'data' dimension added, or original if unchanged
-  
+
   Raises:
     AssertionError: If sharding is not NamedSharding or shape cannot be sharded
   """
@@ -1000,10 +1004,12 @@ def add_data_to_sharding(mesh, path, aval, sharding):
   try:
     sharded_shape = sharding.shard_shape(aval.shape)
   except Exception as e:
-    raise AssertionError(f"Could not shard value {jax.tree_util.keystr(path)} of shape={aval.shape} with {sharding=}") from e
+    raise AssertionError(
+        f"Could not shard value {jax.tree_util.keystr(path)} of shape={aval.shape} with {sharding=}"
+    ) from e
   pspec = sharding.spec
 
-  if 'data' in jax.tree.leaves(pspec):
+  if "data" in jax.tree.leaves(pspec):
     return sharding
 
   for idx, (size, partition) in enumerate(zip(sharded_shape, pspec)):
@@ -1013,25 +1019,26 @@ def add_data_to_sharding(mesh, path, aval, sharding):
     if isinstance(partition, str):
       partition = (partition,)
 
-    if size % mesh.shape['data'] == 0 and (partition is None or 'tensor' not in partition):
-      added_component = ('data',) + partition
-      new_pspec = jax.sharding.PartitionSpec(*(pspec[:idx] + (added_component,) + pspec[idx+1:]))
+    if size % mesh.shape["data"] == 0 and (partition is None or "tensor" not in partition):
+      added_component = ("data",) + partition
+      new_pspec = jax.sharding.PartitionSpec(*(pspec[:idx] + (added_component,) + pspec[idx + 1 :]))
       new_sharding = jax.sharding.NamedSharding(sharding.mesh, new_pspec)
       return new_sharding
   return sharding
 
+
 def maybe_update_params_sharding_with_opt(config, state_mesh_shardings):
   """Updates parameter sharding configuration when optimizer state sharding is enabled.
-  
+
   When shard_optimizer_over_data is enabled (Zero-1 style sharding), this function
   extracts the optimizer state shardings from the Adam optimizer's first moment (mu)
   and merges them with the parameter shardings. This ensures parameter sharding is
   consistent with how the optimizer state is distributed across the compute mesh.
-  
+
   Args:
     config: Configuration object with shard_optimizer_over_data flag
     state_mesh_shardings: Train state mesh shardings containing params and opt_state
-  
+
   Returns:
     A tuple of (prev_params_shardings, updated_state_mesh_shardings):
       - prev_params_shardings: Original parameter shardings before the update
@@ -1042,17 +1049,22 @@ def maybe_update_params_sharding_with_opt(config, state_mesh_shardings):
   if config.shard_optimizer_over_data:
     if isinstance(state_mesh_shardings.opt_state, optax.ScaleByAdamState):
       sharded_fp32_params = state_mesh_shardings.opt_state.mu
-    elif isinstance(state_mesh_shardings.opt_state, tuple) and isinstance(state_mesh_shardings.opt_state[0], optax.ScaleByAdamState):
+    elif isinstance(state_mesh_shardings.opt_state, tuple) and isinstance(
+        state_mesh_shardings.opt_state[0], optax.ScaleByAdamState
+    ):
       sharded_fp32_params = state_mesh_shardings.opt_state[0].mu
     else:
-      raise NotImplementedError(f"Could not find optimizer state shardings from optimizer of type {type(state_mesh_shardings.opt_state)}")
+      raise NotImplementedError(
+          f"Could not find optimizer state shardings from optimizer of type {type(state_mesh_shardings.opt_state)}"
+      )
     if "params" not in sharded_fp32_params.keys():
       # When quantization=fp8 is enabled the sharded_fp32_params
       # are not wrapped in `params`. Here we wrap them back.
       sharded_fp32_params = {"params": sharded_fp32_params}
     state_mesh_shardings = state_mesh_shardings.replace(params=dict(prev_params_shardings, **sharded_fp32_params))
   return prev_params_shardings, state_mesh_shardings
 
+
 def setup_initial_state(
     model,
     data_iterator,
@@ -1145,7 +1157,11 @@ def get_abstract_state(model, tx, config, rng, mesh, is_training=True):
   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(add_data_to_sharding, mesh), unbox_logicallypartioned(abstract_state).opt_state, state_mesh_shardings.opt_state)
+        opt_state=jax.tree.map_with_path(
+            functools.partial(add_data_to_sharding, mesh),
+            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)

src/MaxText/train.py

@@ -218,11 +218,13 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
     # When using Zero-1 optimizer sharding, cast params to lower precision and apply sharding constraints
     # so that all-gather is done once in the lower precision before the gradient accumulation loop
     if config.shard_optimizer_over_data:
+
       def convert_to_bf16(param):
         if param.dtype == jnp.float32:
           return param.astype(jnp.bfloat16)
         else:
           return param
+
       ga_params = jax.tree_util.tree_map(convert_to_bf16, params)
       ga_params = jax.tree.map(jax.lax.with_sharding_constraint, ga_params, params_shardings)
     else:
@@ -401,7 +403,9 @@ def train_loop(config, recorder, state=None):
       state = _merge_dpo_state(state, reference_params)
     state_mesh_shardings = _merge_dpo_state(state_mesh_shardings, state_mesh_shardings.params["params"])
 
-  params_shardings, state_mesh_shardings = maxtext_utils.maybe_update_params_sharding_with_opt(config, state_mesh_shardings)
+  params_shardings, state_mesh_shardings = maxtext_utils.maybe_update_params_sharding_with_opt(
+      config, state_mesh_shardings
+  )
 
   p_train_step, p_eval_step = train_utils.jit_train_and_eval_step(
       config, model, mesh, state, state_mesh_shardings, train_step, eval_step, eval_data_iterator, params_shardings

src/MaxText/train_utils.py

@@ -76,9 +76,7 @@ def create_training_tools(config, model, mesh):
   return init_rng, checkpoint_manager, learning_rate_schedule, tx
 
 
-def jit_train_step(
-    config, model, state, state_mesh_shardings, data_sharding, train_step, params_shardings
-):
+def jit_train_step(config, model, state, state_mesh_shardings, data_sharding, train_step, params_shardings):
   """Returns a JIT-compiled train step function, which is loaded from a file if specified in the config."""
   (
       functional_train,
@@ -144,9 +142,7 @@ def jit_train_and_eval_step(
 ):
   """Returns a JIT-compiled train and eval step function."""
   data_sharding = maxtext_utils.get_input_data_sharding(config, mesh)
-  p_train_step = jit_train_step(
-      config, model, state, state_mesh_shardings, data_sharding, train_step, params_shardings
-  )
+  p_train_step = jit_train_step(config, model, state, state_mesh_shardings, data_sharding, train_step, params_shardings)
   p_eval_step = None
   if eval_data_iterator:
     p_eval_step = jit_eval_step(config, model, state_mesh_shardings, data_sharding, eval_step)