diff --git a/src/MaxText/configs/base.yml b/src/MaxText/configs/base.yml
index edb0a661b1..e1a89493a3 100644
--- a/src/MaxText/configs/base.yml
+++ b/src/MaxText/configs/base.yml
@@ -178,6 +178,7 @@ megablox: true
 sparse_matmul: true
 capacity_factor: -1.0 # a factor to decide expert capacity for token dropping, and no dropping by default
 load_balance_loss_weight: 0.01 # weight for the load balance loss
+expert_balance: False # whether or not to do expert balancing
 use_random_routing: false # whether to use random routing for debug/test purpose
 use_custom_sort_vjp: true # whether to use a custom sort vjp for sparse matmul ops
 use_ring_of_experts: false # whether to use ring of experts for sparse matmul expert parallelism
diff --git a/src/MaxText/configs/types.py b/src/MaxText/configs/types.py
index 3ee4428aa8..cb38b0141d 100644
--- a/src/MaxText/configs/types.py
+++ b/src/MaxText/configs/types.py
@@ -551,6 +551,7 @@ class MoEGeneral(BaseModel):
   num_experts_per_tok: PositiveInt = Field(1, description="The number of experts to route each token to.")
   capacity_factor: float = Field(-1.0, description="Expert capacity factor. If < 0, no token dropping.")
   load_balance_loss_weight: NonNegativeFloat = Field(0.01, description="Weight for the load balancing auxiliary loss.")
+  expert_balance: bool = Field(False, description="Whether to use expert balancing.")
   use_custom_sort_vjp: bool = Field(True, description="Whether to use a custom sort VJP for sparse matmul ops.")
   use_ring_of_experts: bool = Field(
       False,
diff --git a/src/MaxText/input_pipeline/_hf_data_processing.py b/src/MaxText/input_pipeline/_hf_data_processing.py
index e056cd972e..82d4641c73 100644
--- a/src/MaxText/input_pipeline/_hf_data_processing.py
+++ b/src/MaxText/input_pipeline/_hf_data_processing.py
@@ -192,6 +192,7 @@ def preprocessing_pipeline(
     use_sft=None,
     sft_train_on_completion_only=True,
     grain_worker_count=1,  # only support 0 or 1
+    max_segments_per_seq = 1, # max segments per sequence
 ):
   """pipeline for preprocessing HF dataset"""
 
@@ -298,10 +299,11 @@ def lists2array(x):
   if packing and not use_dpo:
     length_struct = {col: max_target_length for col in data_column_names}
     operations.append(
-        grain.experimental.PackAndBatchOperation(
-            batch_size=global_batch_size // jax.process_count(),
-            length_struct=length_struct,
-        )
+      grain.experimental.PackAndBatchOperation(
+        batch_size=global_batch_size // jax.process_count(),
+        length_struct=length_struct,
+        max_sequences_per_bin=max_segments_per_seq,
+      )
     )
     operations.append(_input_pipeline_utils.ReformatPacking(data_column_names))
   else:
@@ -386,6 +388,7 @@ def make_hf_train_iterator(
         use_sft=config.use_sft,
         sft_train_on_completion_only=config.sft_train_on_completion_only,
         chat_template_path=config.chat_template_path,
+        max_segments_per_seq=config.max_segments_per_seq,
     )
   return train_iter
 
@@ -437,5 +440,6 @@ def make_hf_eval_iterator(
         use_sft=config.use_sft,
         sft_train_on_completion_only=config.sft_train_on_completion_only,
         chat_template_path=config.chat_template_path,
+        max_segments_per_seq=config.max_segments_per_seq,
     )
   return eval_iter
diff --git a/src/MaxText/layers/attention_op.py b/src/MaxText/layers/attention_op.py
index d136f975cd..6d283c1186 100644
--- a/src/MaxText/layers/attention_op.py
+++ b/src/MaxText/layers/attention_op.py
@@ -1393,7 +1393,7 @@ def cudnn_flash_attention(
       dummy_segment_ids = jnp.ones(shape=query.shape[:2], dtype=jnp.int32)
       dummy_attn_mask = SequenceDescriptor.from_segment_ids_and_pos(segment_ids=dummy_segment_ids, segment_pos=None)
       max_segments_per_seq = self.config.max_segments_per_seq
-    elif using_context_parallelism:
+    elif using_context_parallelism or self.config.dataset_type == "synthetic":
       if self.attention_type == AttentionType.LOCAL_SLIDING:
         raise AssertionError("Sliding window attention is not supported for context parallelism")
       # Context parallelism without packing: only supports causal masking
@@ -1422,7 +1422,7 @@ def cudnn_flash_attention(
         window_size=sliding_window_size,
         context_parallel_causal_load_balanced=self.config.context_parallel_load_balance,
         context_parallel_axis="context",
-        context_parallel_strategy=self.config.context_parallel_strategy,
+        # context_parallel_strategy=self.config.context_parallel_strategy,
         max_segments_per_seq=max_segments_per_seq,
     )
 
diff --git a/src/MaxText/layers/moe.py b/src/MaxText/layers/moe.py
index 51f98d8de8..e1f02217a5 100644
--- a/src/MaxText/layers/moe.py
+++ b/src/MaxText/layers/moe.py
@@ -1530,7 +1530,7 @@ def get_einsum(
       def aqt_einsum(*args, **kwargs):  # pylint: disable=unused-argument
         # simply skip kwargs, since aqt einsum doesn't support any kwargs
         # like precision
-        is_aqt = not isinstance(self.quant, quantizations.Fp8Quantization)
+        is_aqt = not ( isinstance(self.quant, quantizations.Fp8Quantization) or isinstance(self.quant, quantizations.NANOOFp8Quantization) )
         kw = {"mesh_axes": rhs_mesh_axes} if is_aqt else {"dtype": self.dtype}
         return self.quant.einsum(**kw)(*args)  # pytype: disable=attribute-error
 
@@ -1565,6 +1565,27 @@ def dense_matmul(
       wo_bias,
   ) -> tuple[jax.Array, Optional[jax.Array]]:
     """Dense matrix multiplication."""
+    if self.config.expert_balance:
+      ######################################################################################################
+      ############################## start hard code for uniform expert ####################################
+      # Create deterministic rotational pattern for gate logits
+      batch_size, seq_len, num_experts = gate_logits.shape
+
+      # Create base weights for experts (increasing values)
+      base_weights = jnp.linspace(0.1, 0.1 * num_experts, num_experts, dtype=gate_logits.dtype)
+
+      # Create position-based indices matrix [seq_len, num_experts]
+      # Each row represents which index in base_weights to use after rotation
+      indices = (jnp.arange(num_experts)[None, :] + jnp.arange(seq_len)[:, None]) % num_experts
+
+      # Use advanced indexing to create the rotated weights matrix in one operation
+      # This takes the appropriate weight for each position based on the rotation pattern
+      rotated_weights = base_weights[indices]
+
+      # Broadcast to batch dimension
+      gate_logits = jnp.broadcast_to(rotated_weights[None, :, :], (batch_size, seq_len, num_experts))
+      ############################################# end ####################################################
+      ######################################################################################################
     # gate_logits: batch, length, expert
     gate_logits = self._maybe_shard_with_logical(gate_logits, ("activation_batch", "activation_norm_length", None))
     if self.config.model_name.startswith("deepseek3"):
diff --git a/src/MaxText/layers/quantizations.py b/src/MaxText/layers/quantizations.py
index d0f9353b6c..22095fa960 100644
--- a/src/MaxText/layers/quantizations.py
+++ b/src/MaxText/layers/quantizations.py
@@ -296,6 +296,9 @@ def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()):
     """Returns dot_general configured with aqt params."""
     return nn.NANOOFp8DotGeneralOp
 
+  def einsum(self, dtype: DType = jnp.float32):
+    return Fp8Einsum(dtype=dtype,e4m3_dtype=jnp.float8_e4m3fnuz,e5m2_dtype=jnp.float8_e5m2fnuz)
+
 
 def _get_int8_quant_config(config):
   drhs_bits = None
diff --git a/src/MaxText/max_utils.py b/src/MaxText/max_utils.py
index 5b5811e6ba..3f6261620e 100644
--- a/src/MaxText/max_utils.py
+++ b/src/MaxText/max_utils.py
@@ -19,6 +19,7 @@
 from collections.abc import Sequence
 import functools
 from functools import partial
+import json
 import os
 import socket
 import subprocess
@@ -705,6 +706,36 @@ def print_system_information():
   max_logging.log(f"System Information: Jaxlib Version: {jax.lib.__version__}")
   max_logging.log(f"System Information: Jax Backend: {jax.extend.backend.get_backend().platform_version}")
 
+  devices = jax.devices()
+  max_logging.log(f"System Information: Number of devices: {len(devices)}, jax path {jax.__file__}")
+  for i, device in enumerate(devices):
+    if device.local_hardware_id is not None:
+      max_logging.log(
+          f"System Information: Device {i}: {device.id} "
+          f"(Local id: {device.local_hardware_id}, Process index: {device.process_index})"
+      )
+
+
+def save_device_information(config):
+  """Convert device information to JSON format."""
+  devices = jax.devices()
+  device_info = {'hostname': socket.gethostname(), 'devices': []}
+
+  for device in devices:
+    if device.local_hardware_id is not None:
+      info = {
+          "id": device.id,
+          "local_hardware_id": device.local_hardware_id,
+          "process_index": device.process_index,
+          "device_kind": device.device_kind,
+          "platform_version": jax.extend.backend.get_backend().platform_version,
+      }
+      device_info['devices'].append(info)
+  # Save to JSON file
+  device_info_path = os.path.join(config.base_output_directory, "device_info.json")
+  with open(device_info_path, "w") as f:
+    json.dump(device_info, f, indent=4)
+
 
 def permute_to_match_maxtext_rope(arr):
   """Permutes the Huggingface Rope to match the MaxText logic."""
diff --git a/src/MaxText/train.py b/src/MaxText/train.py
index f5e8cf377b..a5a4287889 100644
--- a/src/MaxText/train.py
+++ b/src/MaxText/train.py
@@ -438,6 +438,7 @@ def train_loop(config, recorder, state=None):
             if config.shard_optimizer_over_data:
               state = sharding.maybe_shard_with_name(state, state_mesh_shardings, config.shard_mode)
             state, metrics = p_train_step(state, example_batch, nextrng)
+          jax.block_until_ready(state)
 
       step_time_delta = datetime.datetime.now() - last_step_completion
       last_step_completion = datetime.datetime.now()
@@ -514,6 +515,7 @@ def initialize(argv: Sequence[str]) -> tuple[pyconfig.HyperParameters, Any, Any]
   config = pyconfig.initialize(argv)
   max_utils.print_system_information()
   validate_train_config(config)
+  max_utils.save_device_information(config)
   jax.config.update("jax_use_shardy_partitioner", config.shardy)
   # update explicit sharding-supported config
   if config.shard_mode == ShardMode.EXPLICIT: