Fix multihost_utils.process_allgather tiled=True in hf_utils.py

src/maxtext/checkpoint_conversion/utils/hf_utils.py

@@ -31,7 +31,7 @@
 def convert_jax_weight_to_torch(weight: "jax.Array", dtype: None | str = None) -> torch.Tensor:
   expected_dtype = str(weight.dtype) if dtype is None else dtype
   expected_shape = weight.shape
-  weight = multihost_utils.process_allgather(weight)
+  weight = multihost_utils.process_allgather(weight, tiled=True)
   weight = np.array(weight, dtype="float32")
   torch_dtype = getattr(torch, expected_dtype)
   torch_array = torch.from_numpy(weight).to(torch_dtype).reshape(expected_shape)

src/maxtext/common/common_types.py

@@ -136,7 +136,6 @@ class AttentionType(enum.Enum):
   MLA = "mla"
   COMPRESSED = "compressed"
   FULL = "full"
-  BLOCK_DIFFUSION = "block_diffusion"
 
 
 class ShardMode(enum.Enum):

src/maxtext/configs/base.yml

@@ -388,15 +388,12 @@ param_scan_axis: 1
 # The attention parameter dictates the specific algorithm/methodology used to compute the attention scores
 # The attention_type parameter determines the variants of attention, e.g. global or local_sliding
 attention: 'autoselected' # Supported attention: autoselected, dot_product, flash, cudnn_flash_te
-attention_type: 'global' # Supported attention_type: global, local_sliding, chunk, mla, full, compressed, block_diffusion
+attention_type: 'global' # Supported attention_type: global, local_sliding, chunk, mla
 share_kv_projections: false # Note: Not compatible with attention_type='mla'
 attention_bias: false # If true, adds a learnable bias to the query, key, and value projections
 attention_sink: false
 sliding_window_size: 0
 chunk_attn_window_size: 0
-# Token block size B for block-causal attention in Block Diffusion (arXiv:2503.09573).
-# Tokens attend bidirectionally within a block and causally across blocks.
-causal_block_size: 32
 attn_logits_soft_cap: 0.0
 final_logits_soft_cap: 0.0
 z_loss_multiplier: 0.0
@@ -1119,7 +1116,6 @@ sa_block_kv_dkv_compute: 512
 sa_block_q_dq: 512
 sa_block_kv_dq: 512
 sa_use_fused_bwd_kernel: false
-sa_bwd_dkv_megacore: false  # megacore-parallel kv-head groups in the static dkv grid
 sa_q_layout: "HEAD_DIM_MINOR"
 sa_k_layout: "HEAD_DIM_MINOR"
 sa_v_layout: "HEAD_DIM_MINOR"
@@ -1354,3 +1350,4 @@ elastic_timeout_seconds: 300
 elastic_max_retries: 10
 elastic_min_slice_count: -1
 
+

src/maxtext/configs/pyconfig.py

@@ -38,7 +38,6 @@
 from maxtext.utils import max_logging
 
 logger = logging.getLogger(__name__)
-logger.setLevel(os.environ.get("LOGLEVEL", "INFO"))
 
 _BASE_CONFIG_ATTR = "base_config"
 _MAX_PREFIX = "M_"

src/maxtext/configs/types.py

@@ -582,7 +582,7 @@ class Attention(BaseModel):
       "autoselected",
       description="The attention algorithm to use (dot_product, flash, cudnn_flash_te, vllm_rpa, vllm_batched_rpa, etc).",
   )
-  attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed", "block_diffusion"] = Field(
+  attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed"] = Field(
       "global", description="The variant of attention to use."
   )
   share_kv_projections: bool = Field(
@@ -629,10 +629,6 @@ class Attention(BaseModel):
   )
   sliding_window_size: NonNegativeInt = Field(0, description="The size of the sliding window for local attention.")
   chunk_attn_window_size: NonNegativeInt = Field(0, description="The window size for chunked attention.")
-  causal_block_size: PositiveInt = Field(
-      32,
-      description="The number of token positions in each bidirectional block for block-causal attention.",
-  )
   attn_logits_soft_cap: None | NonNegativeFloat = Field(
       None, description="Soft-cap value for attention logits. None means no cap."
   )
@@ -747,10 +743,6 @@ class SplashAttention(BaseModel):
   sa_block_q_dq: int = Field(512, description="Block size for Q_dq in splash attention.")
   sa_block_kv_dq: int = Field(512, description="Block size for KV_dq in splash attention.")
   sa_use_fused_bwd_kernel: bool = Field(False, description="Use fused backward kernel in splash attention.")
-  sa_bwd_dkv_megacore: bool = Field(
-      False,
-      description="Megacore-parallel kv-head groups in the static dkv grid. Needs >1 KV head; useful at local batch 1.",
-  )
   sa_q_layout: str = Field("HEAD_DIM_MINOR", description="Layout for Q in splash attention.")
   sa_k_layout: str = Field("HEAD_DIM_MINOR", description="Layout for K in splash attention.")
   sa_v_layout: str = Field("HEAD_DIM_MINOR", description="Layout for V in splash attention.")
@@ -3511,19 +3503,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         not isinstance(self.sliding_window_size, int) or self.sliding_window_size <= 0
     ):
       raise ValueError("`sliding_window_size` must be an integer > 0 for 'local_sliding' attention.")
-    if self.attention_type == AttentionType.BLOCK_DIFFUSION.value:
-      if self.packing:
-        # Document-local block origins inside a packed sequence are not tracked
-        # in attention metadata; packing without realignment would cause
-        # cross-document block-attention leakage.
-        raise ValueError("Block-diffusion attention does not support packing; set `packing=False`.")
-      if self.attention not in ("autoselected", "dot_product", "flash"):
-        raise ValueError("Block-diffusion attention is supported only by dot_product attention and TPU Splash attention.")
-      if self.attention in ("autoselected", "flash") and self.hardware != "tpu":
-        raise ValueError(
-            "Block-diffusion attention with attention='autoselected' or attention='flash' requires hardware='tpu'; "
-            "use attention='dot_product' on other hardware."
-        )
     if self.quantize_kvcache and not self.kv_quant_axis:
       raise ValueError("`kv_quant_axis` cannot be empty when quantize_kvcache is True.")
     if (
@@ -3985,32 +3964,9 @@ class RLConfig(
     Decoding,
     IciParallelism,
     DcnParallelism,
-    PipelineParallelism,
-    DilocoParams,
     HardwareAndMesh,
     ModelArchitecture,
-    MTP,
     MoBa,
-    # Advanced Architectures, Tuning, and Optimizers
-    Muon,
-    FineTuning,
-    Distillation,
-    # Datasets and Loading Compatibility
-    DatasetGeneral,
-    TfdsDataset,
-    HfDataset,
-    GrainDataset,
-    OlmoGrainDataset,
-    # Inference, Checkpointing, and Monitoring
-    EmergencyCheckpointing,
-    ElasticTraining,
-    InferenceServer,
-    InferenceBenchmark,
-    PrefixCaching,
-    HloDump,
-    Goodput,
-    GcpMonitoring,
-    ManagedMLDiagnostics,
     # Positional Embeddings
     PositionalEmbedding,
     Rope,
@@ -4037,12 +3993,9 @@ class RLConfig(
     AttentionIndexer,
     SplashAttention,
     Qwen3Next,
-    # Debugging, Profiling, and Telemetry
-    AOT,
+    # Debugging and Profiling
     DevelopmentAndDebugging,
     Profiling,
-    Metrics,
-    Tensorboard,
     # For compatibility with trainer in post_train/rl
     RL,
     RLCluster,
@@ -4051,8 +4004,6 @@ class RLConfig(
     RLReward,
     RLSpecialTokens,
     VLLM,
-    TrainingLoop,
-    DerivedValues,
 ):
   """
   Configuration for Reinforcement Learning in MaxText.
@@ -4238,20 +4189,19 @@ def set_derived_values_and_validate(self) -> "RLConfig":
 
     # Dynamically inject model dimensions.
     emb_scale, num_head_scale, mlp_dim_scale, layer_scale = get_individual_scales(self.global_parameter_scale)
-    self.emb_dim = int((2**emb_scale) * self.base_emb_dim)
-    self.num_query_heads = int((2**num_head_scale) * self.base_num_query_heads)
-    self.num_kv_heads = int((2**num_head_scale) * self.base_num_kv_heads)
-    self.mlp_dim = int((2**mlp_dim_scale) * self.base_mlp_dim)
-    self.moe_mlp_dim = int((2**mlp_dim_scale) * getattr(self, "base_moe_mlp_dim", 0))
-    self.num_decoder_layers = int((2**layer_scale) * self.base_num_decoder_layers)
+    object.__setattr__(self, "emb_dim", int((2**emb_scale) * self.base_emb_dim))
+    object.__setattr__(self, "num_query_heads", int((2**num_head_scale) * self.base_num_query_heads))
+    object.__setattr__(self, "num_kv_heads", int((2**num_head_scale) * self.base_num_kv_heads))
+    object.__setattr__(self, "mlp_dim", int((2**mlp_dim_scale) * self.base_mlp_dim))
+    object.__setattr__(self, "moe_mlp_dim", int((2**mlp_dim_scale) * getattr(self, "base_moe_mlp_dim", 0)))
+    object.__setattr__(self, "num_decoder_layers", int((2**layer_scale) * self.base_num_decoder_layers))
 
     # Mirror into internal MaxText fields for backward compatibility.
     train_micro_batch_size = getattr(self.dataset, "train_micro_batch_size", -1)
     batch_size = getattr(self.dataset, "batch_size", 1)
     if train_micro_batch_size <= 0:
       train_micro_batch_size = batch_size
-    self.micro_batch_size_to_train_on = train_micro_batch_size
-    self.steps = getattr(self, "train_steps", getattr(self, "num_batches", 10))
+    object.__setattr__(self, "micro_batch_size_to_train_on", train_micro_batch_size)
 
     if self.remat_policy == "custom":
       tensors = [
@@ -4276,7 +4226,7 @@ def set_derived_values_and_validate(self) -> "RLConfig":
           "attention_out",
           "out_proj",
       ]
-      self.tensors_on_device = [t for t in tensors if getattr(self, t) == "device"]
-      self.tensors_to_offload = [t for t in tensors if getattr(self, t) == "offload"]
+      object.__setattr__(self, "tensors_on_device", [t for t in tensors if getattr(self, t) == "device"])
+      object.__setattr__(self, "tensors_to_offload", [t for t in tensors if getattr(self, t) == "offload"])
 
     return self

src/maxtext/experimental/agent/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023–2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/adk_agent.py

@@ -11,7 +11,7 @@
 logger = logging.getLogger(__name__)
 
 
-def _send_message_with_retry(chat, prompt, max_retries=5, sleep_seconds=30):
+def _send_message_with_retry(chat, prompt, max_retries=5, sleep_seconds=60):
   """Sends a message to Gemini with retry and a 30-second sleep on 429 rate-limit/quota errors."""
   for attempt in range(1, max_retries + 1):
     try:

src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/mock_failure_log.txt

@@ -0,0 +1,7 @@
+Traceback (most recent call last):
+  File "train.py", line 42, in <module>
+    import maxtext
+  File "/usr/local/google/home/fiyinbenstowe/Desktop/Project/maxtext/src/maxtext/layers/normalizations.py", line 72
+    mean2 = jnp.mean(lax.square(x), axis=-1, keepdims=True)
+          ^
+SyntaxError: invalid syntax

src/maxtext/experimental/agent/ckpt_validation_pipeline/send_email.py

@@ -4,7 +4,7 @@
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#     http://www.apache.org/licenses/LICENSE-;
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
@@ -17,43 +17,54 @@
 # pylint: disable=logging-fstring-interpolation
 
 import argparse
-import base64
-import json
 import os
+import smtplib
 import sys
+from email.message import EmailMessage
 
 from maxtext.utils import max_logging as logger
-from google.cloud import pubsub_v1
 
-def send_alert(subject: str, body: str, recipient: str, attachment_path: str = None):
-  """Dispatches an email alert by publishing it to an Application Integration Pub/Sub topic."""
-
-  project_id = os.environ.get("GOOGLE_CLOUD_PROJECT", "tpu-prod-env-multipod")
-  topic_id = "maxtext-validation-agent-alerts"
 
-  payload = {
-      "subject": subject,
-      "body": body,
-      "recipient": recipient,
-  }
+def send_alert(subject: str, body: str, recipient: str, attachment_path: str = None):
+  """Dispatches an email alert, gracefully degrading to local logging if SMTP is unavailable."""
+  # In a production environment, you would pull these from a secure secrets manager or env.
+  smtp_server = os.environ.get("SMTP_SERVER", "localhost")
+  smtp_port = int(os.environ.get("SMTP_PORT", 1025))  # Default to a local mock server port
+  smtp_user = os.environ.get("SMTP_USERNAME", "")
+  smtp_pass = os.environ.get("SMTP_PASSWORD", "")
+  sender_email = os.environ.get("SENDER_EMAIL", "overwatch-agent@ml-auto-solutions.com")
 
+  msg = EmailMessage()
+  msg.set_content(body)
+  msg["Subject"] = subject
+  msg["From"] = sender_email
+  msg["To"] = recipient
+
   if attachment_path and os.path.exists(attachment_path):
-    with open(attachment_path, "r", encoding="utf-8") as f:
-      payload["attachment_content"] = f.read()
-      payload["attachment_filename"] = os.path.basename(attachment_path)
-
-  data = json.dumps(payload).encode("utf-8")
+    with open(attachment_path, 'rb') as f:
+      file_data = f.read()
+      file_name = os.path.basename(attachment_path)
+    msg.add_attachment(file_data, maintype='text', subtype='markdown', filename=file_name)
 
   try:
-    publisher = pubsub_v1.PublisherClient()
-    topic_path = publisher.topic_path(project_id, topic_id)
-
-    future = publisher.publish(topic_path, data)
-    message_id = future.result(timeout=10)
-    logger.info(f"Published alert to Pub/Sub topic {topic_path}. Message ID: {message_id}")
+    with smtplib.SMTP(smtp_server, smtp_port) as server:
+      server.starttls()
+      if smtp_user and smtp_pass:
+        server.login(smtp_user, smtp_pass)
+      server.send_message(msg)
+    logger.info(f"Successfully sent email alert to {recipient} regarding: {subject}")
+  except ConnectionRefusedError:
+    logger.info(f"Connection to SMTP server {smtp_server}:{smtp_port} refused.")
+    logger.info("Operating in MOCK/DEV mode. The following email WOULD have been sent:")
+    logger.info("--- EMAIL START ---")
+    logger.info(f"To: {recipient}")
+    logger.info(f"Subject: {subject}")
+    logger.info(f"Body:\n{body}")
+    if attachment_path:
+      logger.info(f"Attachment: {attachment_path}")
+    logger.info("--- EMAIL END ---")
   except Exception as e:  # pylint: disable=broad-exception-caught
-    logger.error(f"Failed to push alert to Pub/Sub topic {topic_id}. Is the Integrations API reachable?")
-    logger.error(f"Exception: {e}")
+    logger.error(f"Failed to dispatch email. Exception: {e}")
     sys.exit(1)
 
 

src/maxtext/kernels/attention/tokamax_ring_attention.py

@@ -278,7 +278,6 @@ def build_splash_config(
       dq_reduction_steps=dq_reduction_steps if dq_reduction_steps > 0 else None,
       use_experimental_scheduler=config.use_splash_scheduler,
       ring_scan_unroll=config.ring_scan_unroll,
-      bwd_dkv_megacore=config.sa_bwd_dkv_megacore,
   )
 
 

src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py

@@ -780,7 +780,7 @@ def make_ring_attention(
         mask,
         (bq_dkv, bkv_dkv),
         is_dkv=True,
-        return_dynamic_grid=config.dq_reduction_steps == 3 and not config.bwd_dkv_megacore,
+        return_dynamic_grid=config.dq_reduction_steps == 3,
     )
     assert (mask_function_fwd is None) == (mask_function_dkv is None)
     dkv_mask_sparsity = _mask_sparsity(dkv_mask_info)