minor change to only do bf16 conversion when GA > 1

Add Zero1+GA

fix train_utils.py after rebase

disable shardy when GA>1

add integration test

address comments

fix lint

fix ci parallelism and lint

Author: hx89

src/MaxText/configs/base.yml

@@ -859,3 +859,4 @@ subslice_shape: ""
 
 # NNX
 enable_nnx: false
+shard_optimizer_over_data: False

src/MaxText/layers/decoders.py

@@ -432,7 +432,7 @@ def map_fn(path, value):
       # Apply remat policy to layer
       layer = nn.remat(
           block_layer,
-          prevent_cse=not self.config.scan_layers,
+          prevent_cse=maxtext_utils.should_prevent_cse_in_remat(self.config),
           policy=policy,
           static_argnums=(4, 5),  # Deterministic and model mode are static arguments.
       )

src/MaxText/maxtext_utils.py

@@ -53,9 +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):
+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)
+  functional_train = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings)
   functional_train.__name__ = "train_step"
   in_shardings = (state_mesh_shardings, data_sharding, None)  # State, batch, rng
   out_shardings = (state_mesh_shardings, None)  # State, metrics
@@ -95,6 +95,28 @@ def get_shaped_batch(config):
   return shaped_batch
 
 
+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
+  """
+  if config.scan_layers:
+    return False
+
+  if config.gradient_accumulation_steps > 1 and config.hardware in ("gpu", "gpu_multiprocess"):
+    return False
+
+  return True
+
+
 def load_compiled(config, partial_train, state):
   """# Loading a serialized compiled train step function."""
 
@@ -939,6 +961,97 @@ 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:
+    boxed_pytree: a pytree that includes LogicallyPartitioned
+      leaves.
+  Returns:
+    a pytree where all all LogicallyPartitioned leaves have been unboxed.
+  """
+  return jax.tree_util.tree_map(
+      lambda x: x.unbox() if isinstance(x, nn.spmd.LogicallyPartitioned) else x,
+      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
+  """
+  if not isinstance(sharding, jax.sharding.NamedSharding):
+    raise AssertionError(f"Expected NamedSharding, found {sharding} of {type(sharding)=} at {jax.tree_util.keystr(path)}")
+  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
+  pspec = sharding.spec
+
+  if 'data' in jax.tree.leaves(pspec):
+    return sharding
+
+  for idx, (size, partition) in enumerate(zip(sharded_shape, pspec)):
+    if partition is None:
+      partition = ()
+
+    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:]))
+      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
+      - updated_state_mesh_shardings: State mesh shardings with updated params field
+        (unchanged if shard_optimizer_over_data is False)
+  """
+  prev_params_shardings = state_mesh_shardings.params
+  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):
+      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)}")
+    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,
@@ -1029,6 +1142,11 @@ def get_abstract_state(model, tx, config, rng, mesh, is_training=True):
   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(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)
     state_mesh_shardings = state_mesh_shardings.replace(opt_state=opt_state)

src/MaxText/pyconfig.py

@@ -678,6 +678,21 @@ def user_init(raw_keys):
         1,
     )
 
+    # Automatically disable shardy when gradient accumulation is enabled on GPU
+    # This incompatibility is specific to GPU hardware
+    if (
+        raw_keys["gradient_accumulation_steps"] > 1
+        and raw_keys["shardy"]
+        and raw_keys["hardware"] in ("gpu", "gpu_multiprocess")
+    ):
+      max_logging.log(
+          "WARNING: Automatically setting shardy=False because"
+          f" gradient_accumulation_steps={raw_keys['gradient_accumulation_steps']} > 1"
+          f" on hardware={raw_keys['hardware']}."
+          " Shardy is not compatible with gradient accumulation on GPU."
+      )
+      raw_keys["shardy"] = False
+
     if raw_keys["pagedattn_max_pages_per_group"] <= 0:
       raw_keys["pagedattn_max_pages_per_group"] = (
           raw_keys["max_target_length"] + raw_keys["pagedattn_tokens_per_page"] - 1

src/MaxText/train.py

@@ -190,7 +190,7 @@ def loss_fn(model, config, data, dropout_rng, params, is_train=True):
   return loss, aux
 
 
-def train_step(model, config, state_mesh_shardings, state, data, dropout_rng):
+def train_step(model, config, state_mesh_shardings, params_shardings, state, data, dropout_rng):
   """
 
   Args:
@@ -212,12 +212,27 @@ def train_step(model, config, state_mesh_shardings, state, data, dropout_rng):
     extra_dpo_args = [reference_params]
     _loss_fn = dpo_loss_fn
 
+  params = state.params
+
   if config.gradient_accumulation_steps > 1:
+    # 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:
+      ga_params = params
 
     def accumulate_gradient(acc_grad_and_loss, data):
+      ga_params = acc_grad_and_loss["ga_params"]
       grad_func = jax.value_and_grad(_loss_fn, argnums=4, has_aux=True)
       (_, aux), cur_batch_gradient = grad_func(
-          model, config, data, dropout_rng, state.params, *extra_dpo_args, is_train=True
+          model, config, data, dropout_rng, ga_params, *extra_dpo_args, is_train=True
       )
       acc_grad_and_loss["loss"] += aux["total_loss"]
       acc_grad_and_loss["moe_lb_loss"] += aux["moe_lb_loss"]
@@ -235,8 +250,16 @@ def reshape_to_microbatch_accumulations(batch_arr):
       return jnp.reshape(batch_arr, microbatch_shape)
 
     data = jax.tree_util.tree_map(reshape_to_microbatch_accumulations, data)
-    init_grad = jax.tree_util.tree_map(jnp.zeros_like, state.params)
-    init_grad_and_loss = {"loss": 0.0, "grad": init_grad, "total_weights": 0, "moe_lb_loss": 0.0, "mtp_loss": 0.0}
+    init_grad = jax.tree_util.tree_map(jnp.zeros_like, ga_params)
+    init_grad = jax.tree.map(jax.lax.with_sharding_constraint, init_grad, params_shardings)
+    init_grad_and_loss = {
+        "loss": 0.0,
+        "grad": init_grad,
+        "total_weights": 0,
+        "moe_lb_loss": 0.0,
+        "mtp_loss": 0.0,
+        "ga_params": ga_params,
+    }
 
     grad_and_loss, aux = jax.lax.scan(
         accumulate_gradient, init_grad_and_loss, data, length=config.gradient_accumulation_steps
@@ -246,7 +269,10 @@ def reshape_to_microbatch_accumulations(batch_arr):
         + grad_and_loss["moe_lb_loss"] / config.gradient_accumulation_steps
         + grad_and_loss["mtp_loss"] / config.gradient_accumulation_steps
     )
-    raw_grads = jax.tree_util.tree_map(lambda arr: arr / grad_and_loss["total_weights"], grad_and_loss["grad"])
+    raw_grads = grad_and_loss["grad"]
+    if config.shard_optimizer_over_data:
+      raw_grads = jax.tree.map(jax.lax.with_sharding_constraint, raw_grads, params_shardings)
+    raw_grads = jax.tree_util.tree_map(lambda arr: arr / grad_and_loss["total_weights"], raw_grads)
     aux = jax.tree.map(lambda x: jnp.sum(x, axis=0), aux)  # pytype: disable=module-attr
   else:
     if config.optimizer_memory_host_offload:
@@ -255,8 +281,10 @@ def reshape_to_microbatch_accumulations(batch_arr):
             reference_params, max_utils.with_memory_kind(reference_params_sharding, "device")
         )
         extra_dpo_args = [reference_params]
+    if config.shard_optimizer_over_data:
+      params = jax.tree.map(jax.lax.with_sharding_constraint, params, params_shardings)
     grad_func = jax.value_and_grad(_loss_fn, argnums=4, has_aux=True)
-    (loss, aux), raw_grads = grad_func(model, config, data, dropout_rng, state.params, *extra_dpo_args, is_train=True)
+    (loss, aux), raw_grads = grad_func(model, config, data, dropout_rng, params, *extra_dpo_args, is_train=True)
 
   raw_grads = jax.tree_util.tree_map(lambda x: x.astype(config.grad_dtype) if x.dtype == jnp.float32 else x, raw_grads)
   intermediate_outputs = aux["intermediate_outputs"]
@@ -373,12 +401,15 @@ 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)
+
   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
+      config, model, mesh, state, state_mesh_shardings, train_step, eval_step, eval_data_iterator, params_shardings
   )
 
   with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
     shaped_batch = maxtext_utils.get_shaped_batch(config)
+    state = jax.lax.with_sharding_constraint(state, state_mesh_shardings)
     compiled = p_train_step.lower(state, shaped_batch, init_rng).compile()
     compiled_stats = compiled.memory_analysis()
     max_utils.print_compiled_memory_stats(compiled_stats)
@@ -402,6 +433,7 @@ def train_loop(config, recorder, state=None):
         nextrng = jax.jit(jax.random.fold_in)(init_rng, step)
         with maybe_record_goodput(recorder, GoodputEvent.STEP, step):
           with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+            state = jax.lax.with_sharding_constraint(state, state_mesh_shardings)
             state, metrics = p_train_step(state, example_batch, nextrng)
 
       step_time_delta = datetime.datetime.now() - last_step_completion
@@ -474,9 +506,9 @@ def initialize(argv: Sequence[str]) -> tuple[pyconfig.HyperParameters, Any, Any]
   # TODO: mazumdera@ : ensure missing mandatory fields in base.yml are filled in in argv,
   # or fill in here
   config = pyconfig.initialize(argv)
-  jax.config.update("jax_use_shardy_partitioner", config.shardy)
   max_utils.print_system_information()
   validate_train_config(config)
+  jax.config.update("jax_use_shardy_partitioner", config.shardy)
   os.environ["TFDS_DATA_DIR"] = config.dataset_path or ""
   vertex_tensorboard_manager = VertexTensorboardManager()
   if config.use_vertex_tensorboard or os.environ.get("UPLOAD_DATA_TO_TENSORBOARD"):

src/MaxText/train_utils.py

@@ -76,15 +76,19 @@ 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):
+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,
       in_shardings,
       out_shardings,
       static_argnums,
       donate_argnums,
-  ) = maxtext_utils.get_functional_train_with_signature(train_step, data_sharding, state_mesh_shardings, model, config)
+  ) = maxtext_utils.get_functional_train_with_signature(
+      train_step, data_sharding, state_mesh_shardings, model, config, params_shardings
+  )
 
   # Define the compilation of functional_train, either by loading the compiled version or wrapping a new one in a jit
   if config.compiled_trainstep_file != "":
@@ -136,10 +140,13 @@ def jit_train_and_eval_step(
     train_step,
     eval_step=None,
     eval_data_iterator=None,
+    params_shardings=None,
 ):
   """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)
+  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)

tests/integration_tests/train_tests.py

@@ -347,6 +347,36 @@ def test_base_model_shardy_false(self):
   def test_gpu_synthetic_model_ag_once(self):
     train_main(TrainTests.CONFIGS["synthetic"] + ["model_fsdp_ag_once=True"])
 
+  @pytest.mark.integration_test
+  @pytest.mark.gpu_only
+  def test_gpu_zero1_gradient_accumulation(self):
+    os.environ["NVTE_FUSED_ATTN"] = "1"  # Enable fused attention
+    zero1_ga = [  # tests Zero-1 optimizer sharding with gradient accumulation
+        None,
+        os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml"),
+        "base_output_directory=gs://runner-maxtext-logs",
+        "run_name=runner_test",
+        "dataset_path=gs://maxtext-dataset",
+        "steps=10",
+        "enable_checkpointing=False",
+        "enable_goodput_recording=False",
+        "dataset_type=synthetic",
+        "attention=cudnn_flash_te",
+        "remat_policy=minimal",
+        "scan_layers=False",
+        "max_target_length=8192",
+        "per_device_batch_size=2",
+        "ici_data_parallelism=-1",
+        "dcn_data_parallelism=1",
+        "ici_fsdp_parallelism=1",
+        "dcn_fsdp_parallelism=1",
+        "gradient_accumulation_steps=8",
+        "shard_optimizer_over_data=True",
+        "override_model_config=True",
+        rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizer.llama2')}",
+    ]
+    train_main(zero1_ga)
+
 
 if __name__ == "__main__":
   absltest.main()