diff --git a/src/maxtext/checkpoint_conversion/utils/hf_utils.py b/src/maxtext/checkpoint_conversion/utils/hf_utils.py index 1f50b0f430..77cff2fd80 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_utils.py +++ b/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) diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py index 2155218a58..664ba8b388 100644 --- a/src/maxtext/common/common_types.py +++ b/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): diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index e66213ba51..18d82bebf6 100644 --- a/src/maxtext/configs/base.yml +++ b/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 + diff --git a/src/maxtext/configs/pyconfig.py b/src/maxtext/configs/pyconfig.py index e148aad971..8c00bfb369 100644 --- a/src/maxtext/configs/pyconfig.py +++ b/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_" diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 2145eec47d..6cdf72ad70 100644 --- a/src/maxtext/configs/types.py +++ b/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 diff --git a/src/maxtext/experimental/agent/__init__.py b/src/maxtext/experimental/agent/__init__.py index e69de29bb2..2237c9162e 100644 --- a/src/maxtext/experimental/agent/__init__.py +++ b/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. diff --git a/src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/adk_agent.py b/src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/adk_agent.py index 68344cb8c6..a8c1f0701a 100644 --- a/src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/adk_agent.py +++ b/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: diff --git a/src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/mock_failure_log.txt b/src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/mock_failure_log.txt new file mode 100644 index 0000000000..da609f6d82 --- /dev/null +++ b/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 + 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 diff --git a/src/maxtext/experimental/agent/ckpt_validation_pipeline/send_email.py b/src/maxtext/experimental/agent/ckpt_validation_pipeline/send_email.py index 96ef79bd87..0df8736cae 100644 --- a/src/maxtext/experimental/agent/ckpt_validation_pipeline/send_email.py +++ b/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) diff --git a/src/maxtext/kernels/attention/tokamax_ring_attention.py b/src/maxtext/kernels/attention/tokamax_ring_attention.py index d1bb5e8b75..9fd031599a 100644 --- a/src/maxtext/kernels/attention/tokamax_ring_attention.py +++ b/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, ) diff --git a/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py b/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py index de0ce2b042..4e996516e4 100644 --- a/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py +++ b/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) diff --git a/src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py b/src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py index 696f3a3ae6..a0d7107356 100644 --- a/src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py +++ b/src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py @@ -150,8 +150,6 @@ class SplashConfig: # An experimental scheduler that sometimes produces better softmax overlap. use_experimental_scheduler: bool = False ring_scan_unroll: int = 1 - # Use both TensorCores in the dkv backward by splitting the grid over kv-head groups. - bwd_dkv_megacore: bool = False def __post_init__(self): if self.block_kv_compute is None: @@ -1187,7 +1185,6 @@ def _flash_attention_dkv_kernel( mask_function: MaskFunctionType | None, q_heads_per_kv_head: int, config: SplashConfig, - megacore_groups: bool = False, ): del mask_next_ref, active_cols_ref HEAD_DIM_MINOR = QKVLayout.HEAD_DIM_MINOR @@ -1203,23 +1200,16 @@ def _flash_attention_dkv_kernel( should_initialize = bounds_start_ref[grid_idx].astype(jnp.bool_) should_write = bounds_end_ref[grid_idx].astype(jnp.bool_) else: - if megacore_groups: - kv_index, q_head_index_per_kv_head, q_index = ( - pl.program_id(1), - pl.program_id(2), - pl.program_id(3), - ) - else: - kv_index, q_head, q_index = ( - pl.program_id(0), - pl.program_id(1), - pl.program_id(2), - ) - q_head_index_per_kv_head = lax.rem(q_head, q_heads_per_kv_head) if q_heads_per_kv_head > 1 else 0 + kv_index, q_head, q_index = ( + pl.program_id(0), + pl.program_id(1), + pl.program_id(2), + ) grid_idx = (kv_index * q_steps) + q_index should_initialize = q_index == 0 should_write = True if q_steps <= 2 else q_index == q_steps - 1 if q_heads_per_kv_head > 1: + q_head_index_per_kv_head = lax.rem(q_head, q_heads_per_kv_head) should_initialize = jnp.logical_and(should_initialize, q_head_index_per_kv_head == 0) should_write = jnp.logical_and(should_write, q_head_index_per_kv_head == q_heads_per_kv_head - 1) @@ -1417,12 +1407,6 @@ def _splash_attention_bwd_dkv( kv_steps = kv_seq_len // bkv q_steps = q_seq_len // bq q_heads_per_kv_head = num_q_heads // num_kv_heads - if config.bwd_dkv_megacore and (is_mqa or num_kv_heads == 1): - raise ValueError( - f"bwd_dkv_megacore requires more than one KV head; got is_mqa={is_mqa}, " f"num_kv_heads={num_kv_heads}." - ) - # The kv-head group split needs a static grid. - megacore_groups = config.bwd_dkv_megacore and not dynamic_grid if dynamic_grid: @@ -1442,16 +1426,6 @@ def mask_index_map(h, grid_idx, rows_ref, cols_ref, mask_next_ref=None, *_): next_m = to_i32(mask_next_ref[grid_idx]) # pyrefly: ignore[unsupported-operation] return next_m, 0, 0 - elif megacore_groups: - unravel = lambda f: lambda g, j, gi, i, *_: f(g * q_heads_per_kv_head + gi, i, j) - grid = (num_kv_heads, kv_steps, q_heads_per_kv_head, q_steps) - - def mask_index_map(g, j, gi, i, rows_ref, cols_ref, mask_next_ref=None, *_): - del g, gi, rows_ref, cols_ref # Unused. - grid_idx = j * q_steps + i - next_m = to_i32(mask_next_ref[grid_idx]) # pyrefly: ignore[unsupported-operation] - return next_m, 0, 0 - else: unravel = lambda f: lambda j, h, i, *_: f(h, i, j) grid = (kv_steps, num_q_heads, q_steps) @@ -1625,7 +1599,6 @@ def create_dkv_index_map(h, i, j, *_): bkv=bkv, mask_function=mask_function, q_heads_per_kv_head=q_heads_per_kv_head, - megacore_groups=megacore_groups, ) kernel_name = get_kernel_name( @@ -1776,13 +1749,7 @@ def _bwd_cost_estimate( # 2) for kv_seq_len, the splash attention prefetch schedule assumes no # megacore # 3) for q_seq_len, we are reducing over it to compute dkv - # With megacore_groups the kv-head group dimension is independent (each - # group owns its dk/dv/dq blocks), so it is marked parallel. - compiler_params=pltpu.CompilerParams( - dimension_semantics=(("parallel",) + ("arbitrary",) * (len(grid) - 1)) - if megacore_groups - else (("arbitrary",) * len(grid)) - ), + compiler_params=pltpu.CompilerParams(dimension_semantics=("arbitrary",) * len(grid)), name=kernel_name, cost_estimate=cost_estimate, interpret=config.interpret, @@ -2047,7 +2014,7 @@ def _make_splash_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) diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index f6e2246dd9..e26d2d903f 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -84,16 +84,6 @@ dynamic_vector_slice_in_dim = jax.vmap(lax.dynamic_slice_in_dim, in_axes=(None, 0, None, None)) -def _resolve_attention_type(config: Config, attention_type: AttentionType | str | None) -> AttentionType: - configured_attention_type = AttentionType(getattr(config, "attention_type", AttentionType.GLOBAL.value)) - if attention_type is None: - return configured_attention_type - resolved_attention_type = AttentionType(attention_type) - if configured_attention_type == AttentionType.BLOCK_DIFFUSION and resolved_attention_type == AttentionType.GLOBAL: - return configured_attention_type - return resolved_attention_type - - def validate_compute_axis_order(s: AxisIdxes) -> None: valid_compute_axis_order = ((0, 1, 2, 3), (0, 2, 1, 3)) if s not in valid_compute_axis_order: # currently supported compute_axis_order @@ -214,50 +204,6 @@ def __hash__(self): ) -class BlockCausalMask(splash_attention_mask._ComputableMask): # pylint: disable=protected-access,abstract-method - """Lazy mask with bidirectional attention within causal blocks.""" - - causal_block_size: int - - def __init__( - self, - shape: tuple[int, int], - causal_block_size: int, - shard_count: int = 1, - ): - if causal_block_size <= 0: - raise ValueError("causal_block_size must be positive") - self.causal_block_size = causal_block_size - - def block_causal_mask_function(q_ids, kv_ids): - return (q_ids // self.causal_block_size) >= (kv_ids // self.causal_block_size) - - super().__init__( - shape=shape, - mask_function=block_causal_mask_function, - shard_count=shard_count, - ) - - def __eq__(self, other: object): - if not isinstance(other, type(self)): - return NotImplemented - return ( - self.shape == other.shape - and self.causal_block_size == other.causal_block_size - and np.array_equal(self.q_sequence, other.q_sequence) - ) - - def __hash__(self): - return hash( - ( - type(self), - self.shape, - self.causal_block_size, - self.q_sequence.tobytes() if self.q_sequence is not None else None, - ) - ) - - def _generate_chunk_attention_mask(mask_shape: tuple[int, int], chunk_size: int, q_offset: int = 0) -> jax.Array: """Generates an explicit boolean mask for chunked causal attention. @@ -288,32 +234,6 @@ def _generate_chunk_attention_mask(mask_shape: tuple[int, int], chunk_size: int, return chunk_mask -def _generate_block_causal_attention_mask( - mask_shape: tuple[int, int], causal_block_size: int, q_offset: int = 0 -) -> jax.Array: - """Generates a block-causal mask for Block Diffusion Language Models (BD3LMs). - - Implements the block-causal attention pattern (M_BC) described in Arriola et - al., "Block Diffusion: Interpolating Between Autoregressive and Diffusion - Language Models" (https://arxiv.org/abs/2503.09573). - - For block size B = causal_block_size, query index q, and key index k: - - * Within block i (floor(q/B) == floor(k/B)), tokens attend bidirectionally. - * Across blocks (floor(q/B) > floor(k/B)), block i attends causally to every - preceding block, but not to future blocks. - - B = 1 reduces to standard causal attention, while B equal to the sequence - length reduces to full bidirectional attention. - """ - if causal_block_size <= 0: - raise ValueError("causal_block_size must be positive") - - row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0) + q_offset - col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1) - return (row_ids // causal_block_size) >= (col_ids // causal_block_size) - - def _make_block_mask_indices(bidirectional_mask): """Creates block mask identifying segments based on a bidirectional mask. @@ -559,13 +479,7 @@ def __init__( self.dtype = dtype self.quant = quant self.kv_quant = kv_quant - self.attention_type = _resolve_attention_type(self.config, attention_type) - self.causal_block_size = getattr(self.config, "causal_block_size", None) - if self.attention_type == AttentionType.BLOCK_DIFFUSION: - if self.causal_block_size is None or self.causal_block_size <= 0: - raise ValueError("causal_block_size must be positive for block-diffusion attention") - if self.attention_kernel not in ("autoselected", "dot_product", "flash"): - raise ValueError("Block-diffusion attention is supported only by dot_product attention and TPU Splash attention.") + self.attention_type = attention_type # Block sizes are only used by TPU splash attention kernels. Exclude non-splash kernels if self.attention_kernel not in ( "dot_product", @@ -869,7 +783,6 @@ def generate_attention_mask( if model_mode != MODEL_MODE_AUTOREGRESSIVE and self.attention_type not in ( AttentionType.FULL, AttentionType.COMPRESSED, - AttentionType.BLOCK_DIFFUSION, ): if use_segment_positions: causal_mask = (position_col_ids <= position_row_ids)[:, None, None, :, :] @@ -994,22 +907,6 @@ def _align_mask(m, target_ndim): ) output_mask = chunk_mask * output_mask - # For standard token-by-token autoregressive decoding, keep the existing - # causal mask path unchanged. BD3LM generation instead uses block-level - # parallel sampling. - elif self.attention_type == AttentionType.BLOCK_DIFFUSION and model_mode != MODEL_MODE_AUTOREGRESSIVE: - if use_segment_positions: - block_mask = ((position_row_ids // self.causal_block_size) >= (position_col_ids // self.causal_block_size))[ - :, None, None, :, : - ] - else: - block_mask = _generate_block_causal_attention_mask( - mask_shape=(q_seq_len, kv_seq_len), - causal_block_size=self.causal_block_size, - q_offset=next_pos, - )[None, None, None, :, :] - output_mask = block_mask if output_mask is None else jnp.logical_and(output_mask, block_mask) - if bidirectional_mask is not None: image_mask = _make_bidirectional_block_mask(bidirectional_mask) output_mask = output_mask | image_mask[:, None, None, ...] @@ -1276,11 +1173,6 @@ def apply_attention( return out, None, None else: - if self.attention_type == AttentionType.BLOCK_DIFFUSION: - raise ValueError( - "Block-diffusion flash attention is supported only by TPU Splash; " - "use attention='dot_product' on other hardware." - ) if model_mode == MODEL_MODE_AUTOREGRESSIVE: # fallback to dot_product as pallas gpu flash attention doesn't support decode stage return self.apply_attention_dot( @@ -1567,24 +1459,13 @@ def create_sa_config(config, query, key, attn_logits_soft_cap): sa_config = create_sa_config(self.config, query, key, attn_logits_soft_cap) mask_shape = (query.shape[2], key.shape[2]) # (q_seq_len, kv_seq_len) mask_module = tokamax_splash_mask if self.config.use_tokamax_splash else splash_attention_mask - use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel if self.attention_type == AttentionType.FULL: mask = mask_module.FullMask(mask_shape) - elif self.attention_type == AttentionType.BLOCK_DIFFUSION: - mask_type = LoadBalancedBlockCausalMask if use_load_balanced_cp else BlockCausalMask - mask_kwargs = {"cp_size": cp_size} if use_load_balanced_cp else {} - mask = mask_type( - shape=mask_shape, - causal_block_size=self.causal_block_size, - **mask_kwargs, - ) else: mask = mask_module.CausalMask(shape=mask_shape) - if use_load_balanced_cp and self.attention_type not in ( - AttentionType.FULL, - AttentionType.BLOCK_DIFFUSION, - ): + use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel + if use_load_balanced_cp and self.attention_type != AttentionType.FULL: mask = LoadBalancedCausalMask(shape=mask_shape, cp_size=cp_size) # Apply local masking if local sliding attention is enabled. @@ -2606,21 +2487,3 @@ def __init__( shard_count=shard_count, ) self.q_sequence = _load_balanced_q_sequence(shape, cp_size) - - -class LoadBalancedBlockCausalMask(BlockCausalMask): # pylint: disable=abstract-method - """Lazy block-causal mask with load-balanced query positions.""" - - def __init__( - self, - shape: tuple[int, int], - causal_block_size: int, - cp_size: int, - shard_count: int = 1, - ): - super().__init__( - shape=shape, - causal_block_size=causal_block_size, - shard_count=shard_count, - ) - self.q_sequence = _load_balanced_q_sequence(shape, cp_size) diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index 4314fdb23a..5a65b35ba4 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -51,7 +51,7 @@ AttentionType, ) from maxtext.layers import nnx_wrappers -from maxtext.layers.attention_op import AttentionOp, _resolve_attention_type +from maxtext.layers.attention_op import AttentionOp from maxtext.layers.embeddings import ( LLaMARotaryEmbedding, LlamaVisionRotaryEmbedding, @@ -120,7 +120,7 @@ def attention_as_linen( float32_logits: bool = False, # cast logits in float32 for stability. quant: Optional[Quant] = None, kv_quant: Optional[KVQuant] = None, - attention_type: AttentionType = AttentionType.GLOBAL, + attention_type: AttentionType = AttentionType.GLOBAL, # Default to global attention attn_logits_soft_cap: float | None = None, sliding_window_size: int | None = None, use_ragged_attention: bool = False, @@ -277,7 +277,7 @@ def __init__( float32_logits: bool = False, # cast logits in float32 for stability. quant: Optional[Quant] = None, kv_quant: Optional[KVQuant] = None, - attention_type: AttentionType = AttentionType.GLOBAL, + attention_type: AttentionType = AttentionType.GLOBAL, # Default to global attention attn_logits_soft_cap: float | None = None, sliding_window_size: int | None = None, use_ragged_attention: bool = False, @@ -388,7 +388,7 @@ def __init__( self.float32_logits = float32_logits self.quant = quant self.kv_quant = kv_quant - self.attention_type = _resolve_attention_type(self.config, attention_type) + self.attention_type = attention_type self.attn_logits_soft_cap = attn_logits_soft_cap self.sliding_window_size = sliding_window_size self.use_ragged_attention = use_ragged_attention @@ -851,9 +851,6 @@ def init_rotary_embedding(self): cast_as_fprop_dtype=True, fprop_dtype=self.dtype, mrope_section=self.mrope_section, - partial_rotary_factor=( - self.partial_rotary_factor if self.partial_rotary_factor is not None else self.config.partial_rotary_factor - ), rngs=self.rngs, ) diff --git a/src/maxtext/layers/embeddings.py b/src/maxtext/layers/embeddings.py index 290163f4f6..bae4514ec9 100644 --- a/src/maxtext/layers/embeddings.py +++ b/src/maxtext/layers/embeddings.py @@ -1788,7 +1788,6 @@ def __init__( cast_as_fprop_dtype: bool = True, fprop_dtype: DType = jnp.bfloat16, mrope_section: tuple[int, int, int] | None = None, - partial_rotary_factor: float = 1.0, attention_scaling: float = 1.0, rngs: nnx.Rngs = None, ): @@ -1797,30 +1796,19 @@ def __init__( Args: min_timescale: Start of the geometric index (typically 1). max_timescale: End of the geometric index (rope_theta, e.g., 1000000). - embedding_dims: Dimension of the attention head. + embedding_dims: Dimension of the embedding (head_dim). cast_as_fprop_dtype: Whether to cast output to fprop dtype. fprop_dtype: The dtype of the output. mrope_section: Tuple of (temporal_dim, height_dim, width_dim) for MRoPE. Defaults to [24, 20, 20] if None. - partial_rotary_factor: Fraction of the head dimensions to rotate. The - remaining suffix is passed through unchanged. attention_scaling: Scaling factor applied to cos/sin embeddings. Defaults to 1.0. rngs: rng keys passed in by nnx.bridge.to_linen. """ - if partial_rotary_factor is None or not 0.0 < partial_rotary_factor <= 1.0: - raise ValueError(f"partial_rotary_factor must be in (0, 1], got {partial_rotary_factor}.") - - self.head_dim = embedding_dims - self.partial_rotary_factor = partial_rotary_factor - self.rotary_dim = int(self.head_dim * self.partial_rotary_factor) - if self.rotary_dim <= 0 or self.rotary_dim % 2: - raise ValueError("Rotary dim for rotary position embedding must be a positive multiple of 2.") - super().__init__( min_timescale=min_timescale, max_timescale=max_timescale, mesh=None, - embedding_dims=self.rotary_dim, + embedding_dims=embedding_dims, cast_as_fprop_dtype=cast_as_fprop_dtype, fprop_dtype=fprop_dtype, rngs=rngs, @@ -1828,11 +1816,8 @@ def __init__( self.mrope_section = mrope_section if mrope_section is not None else (24, 20, 20) self.attention_scaling = attention_scaling - if sum(self.mrope_section) != self.rotary_dim // 2: - raise ValueError( - f"mrope_section must describe rotary_dim / 2 frequencies; got {self.mrope_section} " - f"for rotary_dim={self.rotary_dim}." - ) + if self.embedding_dims % 2: + raise ValueError("Embedding dim for rotary position embedding must be a multiple of 2.") def _apply_interleaved_mrope(self, freqs: jax.Array) -> jax.Array: """Apply interleaved MRoPE pattern to 3D rotary embeddings. @@ -1841,16 +1826,16 @@ def _apply_interleaved_mrope(self, freqs: jax.Array) -> jax.Array: interleaved [THTHWHTHW...], preserving frequency continuity. Args: - freqs: Shape (3, batch, seq_len, rotary_dim // 2) + freqs: Shape (3, batch, seq_len, head_dim // 2) Dimension 0: temporal frequencies Dimension 1: height frequencies Dimension 2: width frequencies Returns: - freqs_t: Shape (batch, seq_len, rotary_dim // 2) with interleaved pattern + freqs_t: Shape (batch, seq_len, head_dim // 2) with interleaved pattern """ # Start with temporal frequencies (dimension 0) - freqs_t = freqs[0] # (batch, seq_len, rotary_dim // 2) + freqs_t = freqs[0] # (batch, seq_len, head_dim // 2) # Create interleaved pattern # For each spatial dimension (H, W), place frequencies at positions: @@ -1873,9 +1858,7 @@ def __call__( """Generates rotary position embeddings for multimodal sequences. Args: - inputs: Input tensor of shape [batch, sequence, heads, head_dim]. MRoPE - is applied to the first ``rotary_dim`` features and the rest are - returned unchanged. + inputs: Input tensor of shape [batch, sequence, heads, head_dim]. position: Position IDs with shape: - [batch, sequence] for text-only (2D) - [3, batch, sequence] for multimodal with vision (3D) @@ -1886,8 +1869,10 @@ def __call__( """ if len(inputs.shape) != 4: raise ValueError("Input is assumed to be a rank 4 tensor of shape [batch, sequence, heads, head_dim].") - if self.head_dim != inputs.shape[3]: - raise ValueError("The head dim of the rotary position embedding must match the hidden dimension of the inputs.") + if self.embedding_dims != inputs.shape[3]: + raise ValueError( + "The embedding dims of the rotary position embedding must match the hidden dimension of the inputs." + ) # Handle both 2D (text-only) and 3D (multimodal) position IDs if position.ndim == 2: @@ -1896,27 +1881,25 @@ def __call__( elif position.ndim != 3 or position.shape[0] != 3: raise ValueError(f"Position IDs must be 2D (batch, seq) or 3D (3, batch, seq), got shape {position.shape}") - # Compute frequencies over the rotated prefix only. - inv_freq_expanded = (1.0 / self.timescale)[jnp.newaxis, jnp.newaxis, jnp.newaxis, :] + # Compute frequencies: (3, batch, seq, 1) @ (head_dim // 2, 1) -> (3, batch, seq, head_dim // 2) + inv_freq_expanded = (1.0 / self.timescale)[jnp.newaxis, jnp.newaxis, jnp.newaxis, :] # (1, 1, 1, head_dim//2) position_expanded = position[..., jnp.newaxis] # (3, batch, seq, 1) - freqs = position_expanded * inv_freq_expanded # (3, batch, seq, rotary_dim // 2) + freqs = position_expanded * inv_freq_expanded # (3, batch, seq, head_dim//2) # Apply interleaved MRoPE pattern for 3D positions - freqs = self._apply_interleaved_mrope(freqs) # (batch, seq, rotary_dim // 2) + freqs = self._apply_interleaved_mrope(freqs) # (batch, seq, head_dim//2) # Compute sin and cos - # Duplicate frequencies for the two halves of the rotated prefix. - emb = jnp.concatenate([freqs, freqs], axis=-1) - cos_emb = jnp.cos(emb) * self.attention_scaling - sin_emb = jnp.sin(emb) * self.attention_scaling + # Concatenate to get full head_dim: (batch, seq, head_dim//2) -> (batch, seq, head_dim) + emb = jnp.concatenate([freqs, freqs], axis=-1) # Duplicate for both halves + cos_emb = jnp.cos(emb) * self.attention_scaling # (batch, seq, head_dim) + sin_emb = jnp.sin(emb) * self.attention_scaling # (batch, seq, head_dim) - # Expand for the heads dimension. + # Expand for heads dimension: (batch, seq, head_dim) -> (batch, seq, 1, head_dim) cos_emb = cos_emb[:, :, jnp.newaxis, :] sin_emb = sin_emb[:, :, jnp.newaxis, :] - inputs_rotary, inputs_pass = jnp.split(inputs, [self.rotary_dim], axis=-1) - rotated = self.apply_rotary(inputs_rotary, cos_emb, sin_emb) - x_out = jnp.concatenate([rotated, inputs_pass], axis=-1) + x_out = self.apply_rotary(inputs, cos_emb, sin_emb) if self.cast_as_fprop_dtype: x_out = x_out.astype(self.fprop_dtype) diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py index 584982d1eb..b67420313a 100644 --- a/src/maxtext/trainers/post_train/rl/train_rl.py +++ b/src/maxtext/trainers/post_train/rl/train_rl.py @@ -741,13 +741,9 @@ def _rl_train_impl(argv: Sequence[str], kwargs: dict): # Run evaluation before training if trainer_config.num_test_batches > 0: - # Explicitly sync actor model weights to the rollout engine before Pre-RL evaluation. - # When resuming from an RL checkpoint (step > 0), the trainer restores RL checkpoint - # weights into actor_model after RLCluster initializes. Without this explicit sync, - # the rollout engine would evaluate using the base HuggingFace/SFT weights instead of - # the restored RL checkpoint weights. Calling this unconditionally ensures weight sync - # robustness across all initialization and restore workflows. - rl_cluster.rollout.update_params(nnx.state(actor_model, nnx.Param)) + # `rl_cluster.rollout.update_params()` is intentionally omitted prior to step 0 + # because `RLCluster` initialization (`create_rl_components`) already syncs the actor model weights + # during setup. Skipping this redundant parameter transfer eliminates unnecessary weight resharding. (corr, total, accuracy, partial_accuracy, format_accuracy, mean_reward), _ = evaluate( trainer_config, diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 2a03a6df0d..52ba94a6a9 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -21,7 +21,6 @@ from __future__ import annotations from collections.abc import Callable -import dataclasses from typing import Any from absl import logging @@ -29,17 +28,16 @@ import jax import jax.numpy as jnp from maxtext.common import common_types -from maxtext.common import train_state_nnx from maxtext.configs import pyconfig from maxtext.trainers.pre_train import train as maxtext_train from maxtext.training_engine import abstract_engine from maxtext.training_engine import checkpointing from maxtext.training_engine import inflight_throttler from maxtext.training_engine import metrics as metrics_module +from maxtext.utils import gradient_accumulation from maxtext.utils import max_utils from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils -from maxtext.utils import sharding from maxtext.utils import train_utils @@ -87,7 +85,7 @@ def __init__( self._train_step: int = 0 self._checkpoint_manager = checkpointing.CheckpointManager( - checkpoint_dir=getattr(self._config, "checkpoint_dir", getattr(self._config, "checkpoint_directory", "")), + checkpoint_dir=getattr(self._config, "checkpoint_directory", ""), config=self._config, ) self._metrics_recorder = metrics_module.MetricsRecorder() @@ -102,10 +100,6 @@ def model(self) -> Any: def model(self, new_model: Any) -> None: """Sets the NNX model instance.""" self._model = new_model - self._compiled = False - self._compiled_fwd_bwd = None - self._compiled_update = None - self._model_graphdef = None @property def optimizer(self) -> Any: @@ -116,10 +110,6 @@ def optimizer(self) -> Any: def optimizer(self, new_optimizer: Any) -> None: """Sets the NNX optimizer instance.""" self._optimizer = new_optimizer - self._compiled = False - self._compiled_fwd_bwd = None - self._compiled_update = None - self._state_graphdef = None @property def train_step(self) -> int: @@ -131,32 +121,6 @@ def train_step(self, step: int) -> None: """Sets the current step integer.""" self._train_step = step - @property - def state(self) -> Any: - """Returns the current train state, initializing it if necessary.""" - if self._state is None and self._model is not None and self._optimizer is not None: - self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) - return self._state - - @state.setter - def state(self, new_state: Any) -> None: - """Sets the current train state.""" - self._state = new_state - self._compiled = False - self._compiled_fwd_bwd = None - self._compiled_update = None - self._state_graphdef = None - - @property - def micro_step_count(self) -> int: - """Returns the current micro-batch count in gradient accumulation.""" - return self._micro_step_count - - @property - def has_accumulated_grads(self) -> bool: - """Returns True if accumulated gradients are present.""" - return self._accumulated_grads is not None - def with_loss_fn(self, customized_fn: Callable[..., Any]) -> None: """Overrides the default autoregressive loss function with a custom RL loss. @@ -171,124 +135,12 @@ def with_gen_model_input_fn(self, gen_model_input_fn: Callable[[Any], dict[str, self._gen_model_input_fn = gen_model_input_fn return self - def _fwd_bwd_kernel(self, params, rest, batch): - """Executes a single forward and backward pass to compute gradients.""" - loss_callable = self._loss_fn if self._loss_fn is not None else maxtext_train.loss_fn - - def diff_wrapper(p, r, b): - mdl = nnx.merge(self._model_graphdef, p, r, copy=True) - loss, aux = loss_callable(mdl, self._config, b, None, None, is_train=True) - _, _, new_r = nnx.split(mdl, nnx.Param, ...) - return loss, (aux, new_r) - - grad_func = jax.value_and_grad(diff_wrapper, argnums=0, has_aux=True) - (loss, (aux, new_rest)), micro_grads = grad_func(params, rest, batch) - micro_grads = jax.tree.map( - lambda x: ( - x.astype(getattr(self._config, "grad_dtype", jnp.float32)) - if hasattr(x, "dtype") and x.dtype == jnp.float32 - else x - ), - micro_grads, - ) - return loss, aux, new_rest, micro_grads - - def _update_kernel(self, state_pure, accumulated_grads, micro_step_count, mean_loss): - """Applies accumulated gradients to update the NNX model state.""" - grad_norm = None - is_skipped_val = None - if state_pure is not None: - if micro_step_count <= 1: - grads = accumulated_grads - else: - grads = jax.tree.map( - lambda g: g / micro_step_count, - accumulated_grads, - ) - if getattr(self._config, "gradient_clipping_threshold", 0.0) > 0: - grads = maxtext_utils.apply_gradient_clipping(grads, None, self._config.gradient_clipping_threshold) - local_state = nnx.merge(self._state_graphdef, state_pure, copy=True) - if hasattr(local_state, "apply_gradients"): - if getattr(self._config, "skip_step_on_spikes", False): - grad_norm = max_utils.l2norm_pytree(grads) - local_state.apply_gradients(grads, loss=mean_loss, grad_norm=grad_norm) - opt_obj = getattr(local_state, "optimizer", self._optimizer) - if opt_obj is not None: - opt_state = nnx.to_pure_dict(nnx.state(opt_obj)).get("opt_state", {}) - is_skipped = opt_state.get("is_skipped") if isinstance(opt_state, dict) else None - if is_skipped is not None: - is_skipped_val = is_skipped.astype(jnp.float32) - else: - local_state.apply_gradients(grads) - _, new_state_pure = nnx.split(local_state) - return new_state_pure, grad_norm, is_skipped_val - return state_pure, grad_norm, is_skipped_val - def compile(self, dummy_data: abstract_engine.TrainerPayload) -> None: - """Triggers SPMD JIT compilation of fwd_bwd and update steps. + """Triggers SPMD JIT compilation of fwd_bwd, update, and eval steps. Args: dummy_data: Sample TrainerPayload providing representative tensor shapes. """ - if self._compiled: - return - - if self._state is None: - self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) - - self._state_graphdef, state_pure = nnx.split(self._state) - self._model_graphdef, params_pure, rest_pure = nnx.split(self._model, nnx.Param, ...) - - if self._mesh is not None: - data_sharding = sharding.get_input_data_sharding(self._config, self._mesh) - state_mesh_shardings = jax.tree.map( - lambda x: getattr( - x, - "sharding", - jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec()), - ), - state_pure, - ) - params_shardings = jax.tree.map( - lambda x: getattr( - x, - "sharding", - jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec()), - ), - params_pure, - ) - rest_shardings = jax.tree.map( - lambda x: getattr( - x, - "sharding", - jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec()), - ), - rest_pure, - ) - fwd_bwd_in_shardings = (params_shardings, rest_shardings, data_sharding) - fwd_bwd_out_shardings = (None, None, rest_shardings, params_shardings) - update_in_shardings = (state_mesh_shardings, params_shardings, None) - update_out_shardings = (state_mesh_shardings, None, None) - else: - fwd_bwd_in_shardings = None - fwd_bwd_out_shardings = None - update_in_shardings = None - update_out_shardings = None - - # 1. JIT Compile Micro FWD/BWD Pass - self._compiled_fwd_bwd = jax.jit( - self._fwd_bwd_kernel, - in_shardings=fwd_bwd_in_shardings, - out_shardings=fwd_bwd_out_shardings, - ) - - # 2. JIT Compile Optimizer Update Pass - self._compiled_update = jax.jit( - self._update_kernel, - in_shardings=update_in_shardings, - out_shardings=update_out_shardings, - static_argnums=(2,), - ) self._compiled = True def fwd_bwd(self, payload: abstract_engine.TrainerPayload) -> None: @@ -299,10 +151,9 @@ def fwd_bwd(self, payload: abstract_engine.TrainerPayload) -> None: """ if self._gen_model_input_fn is not None: batch = self._gen_model_input_fn(payload) - elif dataclasses.is_dataclass(payload): - batch = {k: getattr(payload, k) for k in payload.__dataclass_fields__ if getattr(payload, k) is not None} else: batch = payload + loss_callable = self._loss_fn if self._loss_fn is not None else maxtext_train.loss_fn model = getattr(self._state, "model", None) if self._state is not None else self._model if not isinstance(model, nnx.Module): @@ -311,31 +162,28 @@ def fwd_bwd(self, payload: abstract_engine.TrainerPayload) -> None: # Wait for previous computations to finish before dispatching the next one to TPU. self._throttler.wait_for_next() - if self._state is None: - self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) - model = getattr(self._state, "model", self._model) - self._model_graphdef, params, rest = nnx.split(model, nnx.Param, ...) - - if self._compiled and hasattr(self, "_compiled_fwd_bwd"): - loss, aux, new_rest, micro_grads = self._compiled_fwd_bwd(params, rest, batch) - else: - loss, aux, new_rest, micro_grads = self._fwd_bwd_kernel(params, rest, batch) - nnx.update(model, new_rest) + # TODO(mazumdera): This function call should be pre-compiled. + loss, aux, micro_grads = gradient_accumulation.gradient_accumulation_loss_and_grad( + loss_callable, + self._config, + model, + None, + None, + batch, + None, + ) # Don't add metrics to the throttler queue because metrics are logged after # the update step. self._throttler.add_computation(computation=loss, metrics=None) - if loss is not None: - # TODO(mazumdera): This needs to be modified to become - # if isinstance(loss, abstract_engine.WeightedMetric): + if isinstance(loss, abstract_engine.WeightedMetric): self.record_metrics("loss", loss) # Record auxiliary metrics. if isinstance(aux, dict): for key, value in aux.items(): - if value is not None: - self.record_metrics(key, value) + self.record_metrics(key, value) self._cached_losses.append(loss) if self._accumulated_grads is None: @@ -353,32 +201,40 @@ def update(self) -> None: return if self._learning_rate_schedule is not None: - lr = self._learning_rate_schedule(self.train_step) - self.record_metrics("learning_rate", lr) + try: + lr = self._learning_rate_schedule(self.train_step) + self.record_metrics("learning_rate", lr) + except Exception: # pylint: disable=broad-except + pass # Wait for previous computations to finish before dispatching the update step to TPU. self._throttler.wait_for_next() # TODO(mazumdera): The logic below should be pre-compiled. - if self._state is None: - self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) - self._state_graphdef, state_pure = nnx.split(self._state) - - mean_loss = jnp.mean(jnp.array(self._cached_losses)) if self._cached_losses else jnp.array(0.0) - if self._compiled and hasattr(self, "_compiled_update"): - new_state_pure, grad_norm, is_skipped = self._compiled_update( - state_pure, self._accumulated_grads, self._micro_step_count, mean_loss - ) - else: - new_state_pure, grad_norm, is_skipped = self._update_kernel( - state_pure, self._accumulated_grads, self._micro_step_count, mean_loss + if self._state is not None: + # TODO(mazumdera): Figure out how exactly we should normalize the losses + # (if at all). Given that inputs are varying in size, it not correct to + # simply divide by the number of micro-steps. + grads = jax.tree.map( + lambda g: g / max(self._micro_step_count, 1), + self._accumulated_grads, ) - nnx.update(self._state, new_state_pure) - - if grad_norm is not None: - self.record_metrics("gradient_norm", grad_norm) - if is_skipped is not None: - self.record_metrics("step_skipped", is_skipped) + if getattr(self._config, "gradient_clipping_threshold", 0.0) > 0: + grads = maxtext_utils.apply_gradient_clipping(grads, None, self._config.gradient_clipping_threshold) + if hasattr(self._state, "apply_gradients"): + if getattr(self._config, "skip_step_on_spikes", False): + grad_norm = max_utils.l2norm_pytree(grads) + self.record_metrics("gradient_norm", grad_norm) + mean_loss = jnp.mean(jnp.array(self._cached_losses)) if self._cached_losses else jnp.array(0.0) + self._state.apply_gradients(grads, loss=mean_loss, grad_norm=grad_norm) + opt_obj = getattr(self._state, "optimizer", self._optimizer) + if opt_obj is not None: + opt_state = nnx.to_pure_dict(nnx.state(opt_obj)).get("opt_state", {}) + is_skipped = opt_state.get("is_skipped") if isinstance(opt_state, dict) else None + if is_skipped is not None: + self.record_metrics("step_skipped", is_skipped.astype(jnp.float32)) + else: + self._state.apply_gradients(grads) # Add the state to the throttler queue so jax.block_until_ready() waits # for the optimizer update to complete before logging the metrics. @@ -463,35 +319,15 @@ def restore_checkpoint(self, **kwargs: Any) -> Any: checkpoint_state=checkpoint_state, step=step, ) - if restored_step is None: + if not restored_step: return None logging.info("Checkpoint restored from step %d.", restored_step) self.train_step = restored_step if restored_checkpoint_state.accumulated_metrics: - buffers = [] - for b in restored_checkpoint_state.accumulated_metrics: - if isinstance(b, dict): - wms = {} - for k, wm in b.get("weighted_metrics", {}).items(): - if isinstance(wm, dict): - wms[k] = abstract_engine.WeightedMetric(**wm) - else: - wms[k] = wm - buffers.append( - abstract_engine.MetricsBuffer( - id=b.get("id", 0), - mode=b.get("mode", "train"), - weighted_metrics=wms, - scalar_metrics=b.get("scalar_metrics", {}), - aggregation_fns=b.get("aggregation_fns", {}), - ) - ) - else: - buffers.append(b) # pylint: disable-next=protected-access - self._metrics_recorder._metrics_buffer = buffers + self._metrics_recorder._metrics_buffer = restored_checkpoint_state.accumulated_metrics restored_additional_metadata = None if restored_metadata: @@ -524,7 +360,7 @@ def restore_checkpoint(self, **kwargs: Any) -> Any: def record_metrics( self, name: str, - metric: abstract_engine.WeightedMetric | jax.Array | float | int | dict[str, Any], + metric: abstract_engine.WeightedMetric | jax.Array | float | int, aggregation_fn: Callable[[jax.Array], Any] | None = None, ) -> None: """Records a metric into the buffer, appending to JAX arrays. @@ -534,23 +370,12 @@ def record_metrics( metric: The metric to record. aggregation_fn: The aggregation function to apply to the metric. """ - if metric is None: - return - if isinstance(metric, dict): - for sub_k, sub_v in metric.items(): - if sub_v is not None: - self.record_metrics( - f"{name}/{sub_k}" if name else sub_k, - sub_v, - aggregation_fn=aggregation_fn, - ) - else: - self._metrics_recorder.buffer_metrics( - train_step=self.train_step, - name=name, - metric=metric, - aggregation_fn=aggregation_fn, - ) + self._metrics_recorder.buffer_metrics( + train_step=self.train_step, + name=name, + metric=metric, + aggregation_fn=aggregation_fn, + ) def get_metrics(self, clear_cache: bool = True) -> abstract_engine.MetricsBuffer: """Returns accumulated step metrics as an on-device MetricsBuffer. diff --git a/src/maxtext/utils/gradient_accumulation.py b/src/maxtext/utils/gradient_accumulation.py index 106baee07d..35fdf65503 100644 --- a/src/maxtext/utils/gradient_accumulation.py +++ b/src/maxtext/utils/gradient_accumulation.py @@ -176,12 +176,7 @@ def reshape_to_microbatch_accumulations(batch_arr): unreduced_shardings = jax.tree.map(update_sharding_for_unreduced, params_shardings) raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, unreduced_shardings) raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, params_shardings) - divisor = ( - config.gradient_accumulation_steps - if getattr(config, "use_tunix_gradient_accumulation", False) - else grad_and_loss["total_weights"] - ) - raw_grads = jax.tree_util.tree_map(lambda arr: arr / divisor, raw_grads) + 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 if is_nnx: