From 4fa3753d7c4c9a6a3df1cbb44014e57e1c06efdc Mon Sep 17 00:00:00 2001 From: Hana Joo Date: Wed, 1 Jul 2026 16:21:14 -0700 Subject: [PATCH] Automated Code Change PiperOrigin-RevId: 941351298 --- .../utils/hf_model_configs.py | 22 +-- .../utils/param_mapping.py | 138 +++++++++--------- src/maxtext/inference/decode.py | 14 +- .../inference/inference_microbenchmark.py | 4 +- src/maxtext/inference/kvcache.py | 20 +-- src/maxtext/inference/maxengine/maxengine.py | 20 +-- .../input_pipeline/grain_data_processing.py | 4 +- src/maxtext/input_pipeline/grain_tokenizer.py | 10 +- .../input_pipeline/input_pipeline_utils.py | 20 +-- .../input_pipeline/multihost_dataloading.py | 18 +-- src/maxtext/input_pipeline/olmo_data.py | 2 +- .../input_pipeline/packing/prefill_packing.py | 2 +- .../attention/splash_attention_kernel.py | 72 ++++----- src/maxtext/kernels/gather_reduce_sc.py | 26 ++-- src/maxtext/kernels/megablox/backend.py | 4 +- src/maxtext/kernels/megablox/ops.py | 12 +- .../pallas_mosaic_tpu_v2_gmm_kernel.py | 36 ++--- .../pallas_mosaic_tpu_v2_tgmm_kernel.py | 12 +- src/maxtext/kernels/ragged/ragged_gather.py | 12 +- .../kernels/ragged/ragged_gather_reduce.py | 8 +- src/maxtext/layers/nnx_wrappers.py | 14 +- src/maxtext/layers/quantizations.py | 50 +++---- src/maxtext/multimodal/processor.py | 6 +- src/maxtext/multimodal/processor_gemma3.py | 2 +- src/maxtext/multimodal/processor_gemma4.py | 2 +- .../multimodal/processor_qwen3_omni.py | 41 +++--- src/maxtext/multimodal/utils.py | 12 +- src/maxtext/trainers/diloco/diloco.py | 12 +- .../distillation/distillation_utils.py | 10 +- .../post_train/distillation/train_distill.py | 11 +- src/maxtext/trainers/pre_train/train.py | 6 +- .../trainers/pre_train/train_compile.py | 2 +- .../trainers/tokenizer/train_tokenizer.py | 4 +- .../integration/smoke/train_gpu_smoke_test.py | 2 +- .../smoke/train_int8_smoke_test.py | 2 +- .../train_using_ragged_dot_smoke_test.py | 2 +- tests/integration/sparsity_test.py | 2 +- tests/unit/elastic_utils_test.py | 6 +- tests/utils/attention_test_util.py | 4 +- 39 files changed, 324 insertions(+), 322 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py index 8bb262958c..c72afdebc7 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py @@ -19,7 +19,7 @@ import transformers -if transformers.__version__ >= "5.0.0": +if transformers.__version__ >= "5.0.0": # pyrefly: ignore[missing-attribute] from transformers.configuration_utils import PreTrainedConfig as PTConfig # pytype: disable=import-error else: from transformers.configuration_utils import PretrainedConfig as PTConfig @@ -120,7 +120,7 @@ gemma4_31b_dict = gemma4_26b_dict.copy() gemma4_31b_dict["text_config"] = gemma4_26b_dict["text_config"].copy() -gemma4_31b_dict["text_config"].update( +gemma4_31b_dict["text_config"].update( # pyrefly: ignore[no-matching-overload] { "enable_moe_block": False, "hidden_size": 5376, @@ -265,10 +265,10 @@ try: # Will execute successfully if Transformers is updated with Gemma 4 support - gemma4_26b_config = transformers.Gemma4Config(**gemma4_26b_dict) - gemma4_31b_config = transformers.Gemma4Config(**gemma4_31b_dict) - gemma4_e2b_config = transformers.Gemma4Config(**gemma4_e2b_dict) - gemma4_e4b_config = transformers.Gemma4Config(**gemma4_e4b_dict) + gemma4_26b_config = transformers.Gemma4Config(**gemma4_26b_dict) # pyrefly: ignore[missing-attribute] + gemma4_31b_config = transformers.Gemma4Config(**gemma4_31b_dict) # pyrefly: ignore[missing-attribute] + gemma4_e2b_config = transformers.Gemma4Config(**gemma4_e2b_dict) # pyrefly: ignore[missing-attribute] + gemma4_e4b_config = transformers.Gemma4Config(**gemma4_e4b_dict) # pyrefly: ignore[missing-attribute] except AttributeError: # Graceful fallback to raw dict-based PTConfig if Gemma 4 natively is missing gemma4_26b_config = PTConfig(**gemma4_26b_dict) # pytype: disable=wrong-arg-types @@ -1011,7 +1011,7 @@ # TODO(shuningjin): replace with DeepseekV32Config when available in transformers library -class DeepseekV32Config(PTConfig): +class DeepseekV32Config(PTConfig): # pyrefly: ignore[invalid-inheritance] model_type = "deepseek_v32" def __init__(self, **kwargs): @@ -1089,7 +1089,7 @@ def __init__(self, **kwargs): "use_cache": True, "vocab_size": 201088, } -gpt_oss_20b_config = transformers.GptOssConfig(**gpt_oss_20b_dict) +gpt_oss_20b_config = transformers.GptOssConfig(**gpt_oss_20b_dict) # pyrefly: ignore[bad-argument-type] # from https://huggingface.co/openai/gpt-oss-120b/blob/main/config.json # remove mxfp4 quantization_config, since we are using bf16 @@ -1171,7 +1171,7 @@ def __init__(self, **kwargs): "use_cache": True, "vocab_size": 201088, } -gpt_oss_120b_config = transformers.GptOssConfig(**gpt_oss_120b_dict) +gpt_oss_120b_config = transformers.GptOssConfig(**gpt_oss_120b_dict) # pyrefly: ignore[bad-argument-type] qwen3_omni_30b_a3b_config = transformers.Qwen3OmniMoeConfig( @@ -1481,8 +1481,8 @@ def __init__(self, **kwargs): try: # Will execute successfully if Transformers is updated with Qwen3.5 support - qwen3_5_35b_a3b_config = transformers.Qwen3_5MoeConfig(**qwen3_5_35b_a3b_dict) - qwen3_5_397b_a17b_config = transformers.Qwen3_5MoeConfig(**qwen3_5_397b_a17b_dict) + qwen3_5_35b_a3b_config = transformers.Qwen3_5MoeConfig(**qwen3_5_35b_a3b_dict) # pyrefly: ignore[missing-attribute] + qwen3_5_397b_a17b_config = transformers.Qwen3_5MoeConfig(**qwen3_5_397b_a17b_dict) # pyrefly: ignore[missing-attribute] except AttributeError: qwen3_5_35b_a3b_config = PTConfig(**qwen3_5_35b_a3b_dict) # pytype: disable=wrong-arg-types qwen3_5_397b_a17b_config = PTConfig(**qwen3_5_397b_a17b_dict) # pytype: disable=wrong-arg-types diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 80707a2abd..be0c24baea 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -153,7 +153,7 @@ def GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False hf_indices = list(range(block_idx, num_scanned, attention_pattern_length)) for mx, hf in text_params: key = f"params-decoder-layers-layers_{block_idx}-{mx}" - mapping[key] = [f"model.language_model.layers.{i}.{hf}" for i in hf_indices] + mapping[key] = [f"model.language_model.layers.{i}.{hf}" for i in hf_indices] # pyrefly: ignore[bad-assignment] # Remainder layers (unscanned): params-decoder-layers_remainder-layers_{rem_idx}-{param} if num_remaining > 0: @@ -623,7 +623,7 @@ def QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): if scan_layers: # This block handles scanned layers for both dense and MoE models. - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { "params-decoder-layers-pre_self_attention_layer_norm-scale": [ f"model.layers.{i}.input_layernorm.weight" for i in range(n_layers) @@ -663,7 +663,7 @@ def QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): if num_experts > 1: # For scanned MoE, we create a nested list: [[e0_l0, e0_l1..], [e1_l0, e1_l1..]..] # This follows the (experts, layers, ...) tensor layout. - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { "params-decoder-layers-moe_block-gate-kernel": [ f"model.layers.{i}.mlp.gate.weight" for i in range(n_layers) @@ -683,7 +683,7 @@ def QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): } ) else: # Dense MLP - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { "params-decoder-layers-mlp-wi_0-kernel": [ f"model.layers.{i}.mlp.gate_proj.weight" for i in range(n_layers) @@ -714,7 +714,7 @@ def QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): ) if num_experts > 1: # For each unscanned MoE layer, map the MaxText parameter to a 1D list of all expert weights for that layer. - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"params-decoder-layers_{i}-moe_block-gate-kernel": f"model.layers.{i}.mlp.gate.weight", f"params-decoder-layers_{i}-moe_block-wi_0": [ @@ -871,10 +871,10 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals prefix = f"params-decoder-layers-layer_{block_idx}" # Layer norms - mapping[f"{prefix}-input_layernorm-scale"] = [ + mapping[f"{prefix}-input_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] f"model.language_model.layers.{i}.input_layernorm.weight" for i in hf_indices ] - mapping[f"{prefix}-post_attention_layernorm-scale"] = [ + mapping[f"{prefix}-post_attention_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] f"model.language_model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices ] @@ -882,7 +882,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0 if is_full_attention_layer: - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-attention-attention-query-kernel": [ f"model.language_model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices @@ -906,7 +906,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals ) else: # Linear/Hybrid Attention Block - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { # Provide a tuple of HF keys so MaxText concatenates them into qkvz f"{prefix}-attention-in_proj_qkvz-kernel": [ @@ -941,7 +941,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals ) # 3. Handle MLP: Gates and Shared Experts - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-mlp-routed_experts-gate-kernel": [ f"model.language_model.layers.{i}.mlp.gate.weight" for i in hf_indices @@ -962,7 +962,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals ) # 4. Handle MoE Routed Experts - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-mlp-routed_experts-wo": [ f"model.language_model.layers.{i}.mlp.experts.down_proj" for i in hf_indices @@ -999,7 +999,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals ) else: # Linear/Hybrid Attention Block (Unscanned) - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { # Provide a tuple of HF keys so MaxText concatenates them into qkvz f"{prefix}-attention-in_proj_qkvz-kernel": ( @@ -1033,7 +1033,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals ) # MoE Routed Experts - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-mlp-routed_experts-wo": f"model.language_model.layers.{i}.mlp.experts.down_proj", ( @@ -1255,7 +1255,7 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): if is_full_attention_layer: for key in ["query", "key", "value", "out"]: - hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel + hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment] else: hooks[f"{prefix}-attention-in_proj_qkvz-kernel"] = concat_qkvz_and_transpose hooks[f"{prefix}-attention-in_proj_ba-kernel"] = concat_ba_and_transpose @@ -1269,7 +1269,7 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): hooks[f"{mlp_prefix}-shared_expert-wo-kernel"] = transpose hooks[f"{mlp_prefix}-shared_expert_gate-kernel"] = transpose - hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = process_wi_0_wi_1 + hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = process_wi_0_wi_1 # pyrefly: ignore[unsupported-operation] hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose_expert # Vision hooks for Qwen3.5 @@ -1340,23 +1340,23 @@ def reshape_vision_attn_out(input_tensor, target_shape): return input_tensor.T.reshape(target_shape) # Apply vision hooks - hooks["params-vision_encoder-Qwen3_5MoeVisionEncoder_0-patch_embed-proj-kernel"] = reshape_conv3d_patch_embed + hooks["params-vision_encoder-Qwen3_5MoeVisionEncoder_0-patch_embed-proj-kernel"] = reshape_conv3d_patch_embed # pyrefly: ignore[bad-assignment] for i in range(n_vision_layers): prefix = f"params-vision_encoder-Qwen3_5MoeVisionEncoder_0-blocks_{i}" - hooks[f"{prefix}-attn-attn-query-kernel"] = split_qkv_query - hooks[f"{prefix}-attn-attn-query-bias"] = split_qkv_bias_query - hooks[f"{prefix}-attn-attn-key-kernel"] = split_qkv_key - hooks[f"{prefix}-attn-attn-key-bias"] = split_qkv_bias_key - hooks[f"{prefix}-attn-attn-value-kernel"] = split_qkv_value - hooks[f"{prefix}-attn-attn-value-bias"] = split_qkv_bias_value - hooks[f"{prefix}-attn-attn-out-kernel"] = reshape_vision_attn_out - hooks[f"{prefix}-mlp-kernel"] = reshape_kernel_vision - hooks[f"{prefix}-mlp_out-kernel"] = reshape_kernel_vision + hooks[f"{prefix}-attn-attn-query-kernel"] = split_qkv_query # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attn-attn-query-bias"] = split_qkv_bias_query # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attn-attn-key-kernel"] = split_qkv_key # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attn-attn-key-bias"] = split_qkv_bias_key # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attn-attn-value-kernel"] = split_qkv_value # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attn-attn-value-bias"] = split_qkv_bias_value # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attn-attn-out-kernel"] = reshape_vision_attn_out # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-mlp-kernel"] = reshape_kernel_vision # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-mlp_out-kernel"] = reshape_kernel_vision # pyrefly: ignore[bad-assignment] # Vision projector - hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_0-kernel"] = reshape_kernel_vision - hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_2-kernel"] = reshape_kernel_vision + hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_0-kernel"] = reshape_kernel_vision # pyrefly: ignore[bad-assignment] + hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_2-kernel"] = reshape_kernel_vision # pyrefly: ignore[bad-assignment] return hooks @@ -1384,8 +1384,8 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F prefix = f"params-decoder-layers-layer_{block_idx}" # Layer norms - mapping[f"{prefix}-input_layernorm-scale"] = [f"model.layers.{i}.input_layernorm.weight" for i in hf_indices] - mapping[f"{prefix}-post_attention_layernorm-scale"] = [ + mapping[f"{prefix}-input_layernorm-scale"] = [f"model.layers.{i}.input_layernorm.weight" for i in hf_indices] # pyrefly: ignore[bad-assignment] + mapping[f"{prefix}-post_attention_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices ] @@ -1393,7 +1393,7 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0 if is_full_attention_layer: - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-attention-attention-query-kernel": [ f"model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices @@ -1417,7 +1417,7 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F ) else: # Linear/Hybrid Attention Block - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-attention-in_proj_qkvz-kernel": [ f"model.layers.{i}.linear_attn.in_proj_qkvz.weight" for i in hf_indices @@ -1438,7 +1438,7 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F ) # 3. Handle MLP: Gates and Shared Experts - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-mlp-routed_experts-gate-kernel": [f"model.layers.{i}.mlp.gate.weight" for i in hf_indices], f"{prefix}-mlp-shared_expert-wi_0-kernel": [ @@ -1457,7 +1457,7 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F ) # 4. Handle MoE Routed Experts - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-mlp-routed_experts-wi_0": [ [f"model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for i in hf_indices] for e in range(num_experts) @@ -1521,7 +1521,7 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F ) # MoE Routed Experts (List of expert weights for this specific layer) - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-mlp-routed_experts-wi_0": [ f"model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for e in range(num_experts) @@ -1576,7 +1576,7 @@ def permute_conv(input_tensor, target_shape=None): if is_full_attention_layer: for key in ["query", "key", "value", "out"]: - hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel + hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment] else: hooks[f"{prefix}-attention-in_proj_qkvz-kernel"] = transpose hooks[f"{prefix}-attention-in_proj_ba-kernel"] = transpose @@ -1663,17 +1663,17 @@ def DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fal # scan if scan_layers: for maxtext_key, hf_key in dense_layer_keys.items(): - mapping[f"params-decoder-dense_layers-{maxtext_key}"] = [ + mapping[f"params-decoder-dense_layers-{maxtext_key}"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{i}.{hf_key}" for i in range(first_num_dense_layers) ] for maxtext_key, hf_key in moe_layer_keys.items(): - mapping[f"params-decoder-moe_layers-{maxtext_key}"] = [ + mapping[f"params-decoder-moe_layers-{maxtext_key}"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{i}.{hf_key}" for i in range(first_num_dense_layers, num_main_layers) ] for maxtext_key, hf_key in moe_expert_keys.items(): - mapping[f"params-decoder-moe_layers-{maxtext_key}"] = [ + mapping[f"params-decoder-moe_layers-{maxtext_key}"] = [ # pyrefly: ignore[bad-assignment] [f"model.layers.{i}.mlp.experts.{e}.{hf_key}" for i in range(first_num_dense_layers, num_main_layers)] for e in range(num_experts) ] @@ -1690,7 +1690,7 @@ def DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fal mapping[f"params-decoder-moe_layers_{moe_layer_idx}-{maxtext_key}"] = f"model.layers.{i}.{hf_key}" for maxtext_key, hf_key in moe_expert_keys.items(): - mapping[f"params-decoder-moe_layers_{moe_layer_idx}-{maxtext_key}"] = [ + mapping[f"params-decoder-moe_layers_{moe_layer_idx}-{maxtext_key}"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{i}.mlp.experts.{e}.{hf_key}" for e in range(num_experts) ] return mapping @@ -1833,7 +1833,7 @@ def GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals f"model.layers.{i}.mlp.experts.gate_up_proj_bias" for i in hf_indices ], } - mapping.update(block_mapping) + mapping.update(block_mapping) # pyrefly: ignore[no-matching-overload] else: # Unscan @@ -1868,7 +1868,7 @@ def GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals f"{prefix}-GptOssMlp-wi_1_bias", ): f"model.layers.{i}.mlp.experts.gate_up_proj_bias", } - mapping.update(layer_mapping) + mapping.update(layer_mapping) # pyrefly: ignore[no-matching-overload] return mapping @@ -1916,7 +1916,7 @@ def interleave(input_tensor, target_shape=None): """ if saving_to_hf: wi_0, wi_1 = input_tensor - wi_0_1 = np.empty(target_shape, dtype=wi_0.dtype) + wi_0_1 = np.empty(target_shape, dtype=wi_0.dtype) # pyrefly: ignore[no-matching-overload] wi_0_1[..., ::2] = wi_0 wi_0_1[..., 1::2] = wi_1 return wi_0_1 @@ -1936,14 +1936,14 @@ def interleave(input_tensor, target_shape=None): prefix = f"params-decoder-layers-layers_{idx}" if scan_layers else f"params-decoder-layers_{idx}" # Attention Kernels & Biases for key in ["query", "key", "value"]: - hooks[f"{prefix}-GptOssAttention-{key}-kernel"] = reshape_kernel + hooks[f"{prefix}-GptOssAttention-{key}-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment] hooks[f"{prefix}-GptOssAttention-{key}-bias"] = reshape_bias - hooks[f"{prefix}-GptOssAttention-out-kernel"] = reshape_kernel + hooks[f"{prefix}-GptOssAttention-out-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment] # MLP Kernels & Biases hooks[f"{prefix}-GptOssMlp-gate-kernel"] = transpose # `composite_mt_key`: A hook for combining multiple MaxText params. - hooks[(f"{prefix}-GptOssMlp-wi_0", f"{prefix}-GptOssMlp-wi_1")] = interleave - hooks[(f"{prefix}-GptOssMlp-wi_0_bias", f"{prefix}-GptOssMlp-wi_1_bias")] = interleave + hooks[(f"{prefix}-GptOssMlp-wi_0", f"{prefix}-GptOssMlp-wi_1")] = interleave # pyrefly: ignore[unsupported-operation] + hooks[(f"{prefix}-GptOssMlp-wi_0_bias", f"{prefix}-GptOssMlp-wi_1_bias")] = interleave # pyrefly: ignore[unsupported-operation] return hooks @@ -2407,31 +2407,31 @@ def LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals } if scan_layers: - mapping["params-decoder-layers-self_attention-query-kernel"] = [ + mapping["params-decoder-layers-self_attention-query-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.self_attn.q_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-self_attention-key-kernel"] = [ + mapping["params-decoder-layers-self_attention-key-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.self_attn.k_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-self_attention-value-kernel"] = [ + mapping["params-decoder-layers-self_attention-value-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.self_attn.v_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-self_attention-out-kernel"] = [ + mapping["params-decoder-layers-self_attention-out-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.self_attn.o_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-mlp-wi_0-kernel"] = [ + mapping["params-decoder-layers-mlp-wi_0-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.mlp.gate_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-mlp-wi_1-kernel"] = [ + mapping["params-decoder-layers-mlp-wi_1-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.mlp.up_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-mlp-wo-kernel"] = [ + mapping["params-decoder-layers-mlp-wo-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.mlp.down_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-pre_self_attention_layer_norm-scale"] = [ + mapping["params-decoder-layers-pre_self_attention_layer_norm-scale"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.input_layernorm.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-post_self_attention_layer_norm-scale"] = [ + mapping["params-decoder-layers-post_self_attention_layer_norm-scale"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.post_attention_layernorm.weight" for layer_idx in range(n_layers) ] else: @@ -2816,7 +2816,7 @@ def GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False for layer_in_block in range(attention_pattern_length): hf_indices = list(range(layer_in_block, num_scanned, attention_pattern_length)) prefix = f"params-decoder-scanned_blocks-layers_{layer_in_block}" - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-self_attention-query-kernel": [ f"{text_base}.layers.{i}.self_attn.q_proj.weight" for i in hf_indices @@ -2841,14 +2841,14 @@ def GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False } ) if maxtext_config.v_norm_with_scale: - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-self_attention-value_norm-scale": [ f"{text_base}.layers.{i}.self_attn.v_norm.weight" for i in hf_indices ] } ) - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-pre_self_attention_norm-scale": [ f"{text_base}.layers.{i}.input_layernorm.weight" for i in hf_indices @@ -2946,7 +2946,7 @@ def GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False ) if maxtext_config.v_norm_with_scale: mapping.update({f"{prefix}-self_attention-value_norm-scale": (f"{hf_prefix}.self_attn.v_norm.weight")}) - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-pre_self_attention_norm-scale": (f"{hf_prefix}.input_layernorm.weight"), f"{prefix}-post_self_attention_norm-scale": (f"{hf_prefix}.post_attention_layernorm.weight"), @@ -2998,7 +2998,7 @@ def GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False is_global = i % 6 == 5 num_experts = tcfg.get("num_experts") num_experts = num_experts if num_experts is not None else 1 - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-self_attention-query-kernel": (f"{hf_prefix}.self_attn.q_proj.weight"), f"{prefix}-self_attention-key-kernel": (f"{hf_prefix}.self_attn.k_proj.weight"), @@ -3012,7 +3012,7 @@ def GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False ) if maxtext_config.v_norm_with_scale: mapping.update({f"{prefix}-self_attention-value_norm-scale": (f"{hf_prefix}.self_attn.v_norm.weight")}) - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-pre_self_attention_norm-scale": (f"{hf_prefix}.input_layernorm.weight"), f"{prefix}-post_self_attention_norm-scale": (f"{hf_prefix}.post_attention_layernorm.weight"), @@ -3111,7 +3111,7 @@ def GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers } ) if not is_shared: - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-self_attention-key-kernel": f"{hf_prefix}.self_attn.k_proj.weight", f"{prefix}-self_attention-value-kernel": f"{hf_prefix}.self_attn.v_proj.weight", @@ -3527,7 +3527,7 @@ def OLMO3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False) hf_indices = range(cycle_idx, n_layers, layer_cycle_interval) prefix = f"params-decoder-layers-layers_{cycle_idx}" - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { # Attention Projections f"{prefix}-attention-query-kernel": [f"model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices], @@ -3791,9 +3791,9 @@ def process_qkv_vision(input_tensor, target_shape=None): q_hf = input_tensor[:hidden_size, :] k_hf = input_tensor[hidden_size : 2 * hidden_size, :] v_hf = input_tensor[2 * hidden_size :, :] - q_mt = q_hf.T.reshape(target_shape[0]) - k_mt = k_hf.T.reshape(target_shape[1]) - v_mt = v_hf.T.reshape(target_shape[2]) + q_mt = q_hf.T.reshape(target_shape[0]) # pyrefly: ignore[unsupported-operation] + k_mt = k_hf.T.reshape(target_shape[1]) # pyrefly: ignore[unsupported-operation] + v_mt = v_hf.T.reshape(target_shape[2]) # pyrefly: ignore[unsupported-operation] return np.stack([q_mt, k_mt, v_mt], axis=-1) def process_qkv_bias_vision(input_tensor, target_shape=None): @@ -3808,9 +3808,9 @@ def process_qkv_bias_vision(input_tensor, target_shape=None): qb_hf = input_tensor[:hidden_size] kb_hf = input_tensor[hidden_size : 2 * hidden_size] vb_hf = input_tensor[2 * hidden_size :] - qb_mt = qb_hf.reshape(target_shape[0]) - kb_mt = kb_hf.reshape(target_shape[1]) - vb_mt = vb_hf.reshape(target_shape[2]) + qb_mt = qb_hf.reshape(target_shape[0]) # pyrefly: ignore[unsupported-operation] + kb_mt = kb_hf.reshape(target_shape[1]) # pyrefly: ignore[unsupported-operation] + vb_mt = vb_hf.reshape(target_shape[2]) # pyrefly: ignore[unsupported-operation] return np.stack([qb_mt, kb_mt, vb_mt], axis=-1) def reshape_vision_attn_out(input_tensor, target_shape): diff --git a/src/maxtext/inference/decode.py b/src/maxtext/inference/decode.py index 68cd450c99..9b7f415c70 100644 --- a/src/maxtext/inference/decode.py +++ b/src/maxtext/inference/decode.py @@ -139,7 +139,7 @@ def main(argv: Sequence[str]) -> None: if config.use_multimodal: tokens = mm_processor.prepare_text_for_image_fusion(tokens=tokens, config=config, processor_output=processor_outputs) - true_length += image_offsets + true_length += image_offsets # pyrefly: ignore[unbound-name] if config.use_mrope: from maxtext.multimodal import processor_qwen3_omni # pylint: disable=import-outside-toplevel @@ -181,13 +181,13 @@ def main(argv: Sequence[str]) -> None: prefill_result, first_token = engine.prefill( params=params, padded_tokens=tokens, - positions=position_ids, - mrope_deltas=mrope_position_deltas, - images=processor_outputs.pixel_values if config.use_multimodal else None, - image_masks=processor_outputs.pixel_mask if config.use_multimodal and "llama4" in config.model_name else None, + positions=position_ids, # pyrefly: ignore[bad-argument-type] + mrope_deltas=mrope_position_deltas, # pyrefly: ignore[bad-argument-type] + images=processor_outputs.pixel_values if config.use_multimodal else None, # pyrefly: ignore[bad-argument-type] + image_masks=processor_outputs.pixel_mask if config.use_multimodal and "llama4" in config.model_name else None, # pyrefly: ignore[bad-argument-type] videos=getattr(processor_outputs, "video_values", None) if config.use_multimodal else None, - audio_values=processor_outputs.audio_values if config.use_audio else None, - audio_masks=processor_outputs.audio_mask if config.use_audio else None, + audio_values=processor_outputs.audio_values if config.use_audio else None, # pyrefly: ignore[bad-argument-type] + audio_masks=processor_outputs.audio_mask if config.use_audio else None, # pyrefly: ignore[bad-argument-type] true_length=true_length, rng=rng_prefill, slot=i, diff --git a/src/maxtext/inference/inference_microbenchmark.py b/src/maxtext/inference/inference_microbenchmark.py index 1adb0f36d5..d129883af8 100644 --- a/src/maxtext/inference/inference_microbenchmark.py +++ b/src/maxtext/inference/inference_microbenchmark.py @@ -400,8 +400,8 @@ def run_benchmarks(config): config, multisampling_prefill_executable[prefill_length], params, - prefill_tokens[prefill_length], - prefill_true_lengths[prefill_length], + prefill_tokens[prefill_length], # pyrefly: ignore[unbound-name] + prefill_true_lengths[prefill_length], # pyrefly: ignore[unbound-name] benchmark_loop_iters, ) diff --git a/src/maxtext/inference/kvcache.py b/src/maxtext/inference/kvcache.py index ba2266060f..52a3706028 100644 --- a/src/maxtext/inference/kvcache.py +++ b/src/maxtext/inference/kvcache.py @@ -277,7 +277,7 @@ def __init__( *, # Not used in KVCache but passed in by nnx_wrappers.to_linen. # TODO: Remove when bridge no longer needed - rngs: nnx.Rngs = None, + rngs: nnx.Rngs = None, # pyrefly: ignore[bad-function-definition] ): """Initializes the KVCache module. @@ -612,8 +612,8 @@ def kv_cache_chunked_prefill( # For quantized kv cached. Could be get without transpose twice. cached_key = self.get_cached_values(cached_prefill_key_vars, key.dtype, self.prefill_cache_axis_order) cached_value = self.get_cached_values(cached_prefill_value_vars, value.dtype, self.prefill_cache_axis_order) - cached_key_value = jnp.transpose(cached_key, self.prefill_cache_axis_order) - cached_value_value = jnp.transpose(cached_value, self.prefill_cache_axis_order) + cached_key_value = jnp.transpose(cached_key, self.prefill_cache_axis_order) # pyrefly: ignore[bad-argument-type] + cached_value_value = jnp.transpose(cached_value, self.prefill_cache_axis_order) # pyrefly: ignore[bad-argument-type] seq_axis = self.prefill_cache_logical_axis_names.index(CACHE_SEQUENCE) cache_seq_axis = self.prefill_cache_axis_order.index(seq_axis) @@ -767,17 +767,17 @@ def update_ar_key_value( if use_ragged_attention: cache_locations = [slice(None)] * 4 new_token_locations = [slice(None)] * 4 - new_token_locations[ar_cache_sequence_axis] = 0 + new_token_locations[ar_cache_sequence_axis] = 0 # pyrefly: ignore[unsupported-operation] def key_body(i, val): cache_locations[ar_cache_batch_axis] = i - cache_locations[ar_cache_sequence_axis] = lengths[i] + cache_locations[ar_cache_sequence_axis] = lengths[i] # pyrefly: ignore[unsupported-operation] new_token_locations[ar_cache_batch_axis] = i return val.at[tuple(cache_locations)].set(one_token_key_shaped_for_cache[tuple(new_token_locations)]) def value_body(i, val): cache_locations[ar_cache_batch_axis] = i - cache_locations[ar_cache_sequence_axis] = lengths[i] + cache_locations[ar_cache_sequence_axis] = lengths[i] # pyrefly: ignore[unsupported-operation] new_token_locations[ar_cache_batch_axis] = i return val.at[tuple(cache_locations)].set(one_token_value_shaped_for_cache[tuple(new_token_locations)]) @@ -815,7 +815,7 @@ def value_body(i, val): cached_key_scale.set_value( jax.lax.dynamic_update_index_in_dim( cached_key_scale.get_value(), - one_token_key_scale_shaped_for_cache, + one_token_key_scale_shaped_for_cache, # pyrefly: ignore[unbound-name] ar_cache_update_idx, ar_cache_scale_update_axis, ) @@ -823,7 +823,7 @@ def value_body(i, val): cached_value_scale.set_value( jax.lax.dynamic_update_index_in_dim( cached_value_scale.get_value(), - one_token_value_scale_shaped_for_cache, + one_token_value_scale_shaped_for_cache, # pyrefly: ignore[unbound-name] ar_cache_update_idx, ar_cache_scale_update_axis, ) @@ -844,7 +844,7 @@ def get_cached_values(self, cache_vars, target_dtype, cache_axis_order) -> jax.A elif dtype == jnp.float8_e4m3fn: scale_value /= E4M3_MAX - cache_value = KVTensor(qvalue=cache_value, scale=[scale_value], scale_t=None, dequant_dtype=target_dtype, bias=[]) + cache_value = KVTensor(qvalue=cache_value, scale=[scale_value], scale_t=None, dequant_dtype=target_dtype, bias=[]) # pyrefly: ignore[unexpected-keyword] cache_value_in_logical_shape = jax.tree.map(lambda x: reverse_transpose(x, cache_axis_order), cache_value) return cache_value_in_logical_shape @@ -1061,7 +1061,7 @@ def __init__( *, # Not used in MlaKVCache but passed in by nnx_wrappers.to_linen. # TODO: Remove when bridge no longer needed - rngs: nnx.Rngs = None, + rngs: nnx.Rngs = None, # pyrefly: ignore[bad-function-definition] ): """Initializes the MlaKVCache module. diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 5377e86ec3..095d3ea48a 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -101,7 +101,7 @@ def get_keys(self): _BaseEngine = engine_api.Engine if (not is_decoupled() and hasattr(engine_api, "Engine")) else object -class MaxEngine(_BaseEngine): +class MaxEngine(_BaseEngine): # pyrefly: ignore[invalid-inheritance] """The computational core of the generative model server. Engine defines an API that models must adhere to as they plug into the @@ -192,7 +192,7 @@ def _nnx_init_cache_dict(self, mode: str = MODEL_MODE_PREFILL) -> dict: """Zero-filled pure-dict cache matching the abstract NNX model.""" src = self._abstract_model_for_mode(mode) _, cache_state, _ = nnx.split(src, nnx.Cache, ...) - cache_dict = cache_state.to_pure_dict() + cache_dict = cache_state.to_pure_dict() # pyrefly: ignore[missing-attribute] return jax.tree.map(lambda x: jnp.zeros(x.shape, x.dtype), cache_dict) def _nnx_run_model( @@ -219,7 +219,7 @@ def _nnx_run_model( nnx.replace_by_pure_dict(cache_state, cache_dict) # copy=True avoids reusing Variable objects across traces (TraceContextError), # mirroring the workaround in train.py's diff_wrapper. - model = nnx.merge(self.graphdef, params, cache_state, self._nnx_rest_state, copy=True) + model = nnx.merge(self.graphdef, params, cache_state, self._nnx_rest_state, copy=True) # pyrefly: ignore[no-matching-overload] logits = model( decoder_input_tokens, decoder_positions, @@ -280,7 +280,7 @@ def _layout(x, s, l): if x.format == l: return x # Somehow this can be None sometimes. - dll = (l.layout if jax.__version_info__ >= (0, 6, 3) else l.device_local_layout) if isinstance(l, Format) else l + dll = (l.layout if jax.__version_info__ >= (0, 6, 3) else l.device_local_layout) if isinstance(l, Format) else l # pyrefly: ignore[missing-attribute] f = jax.jit(self._identity, out_shardings=Format(dll, s)).lower(x).compile(compiler_options=xla_flags) y = f(x) # Achieves donation of the input argument, but allows for different memory @@ -409,7 +409,7 @@ def _load_params_nnx(self, params, rng): with nn_partitioning.axis_rules(self.config.logical_axis_rules): full_sharding = sharding.nnx_construct_named_sharding(full_abs, self._mesh) concrete_model = maxtext_utils_nnx.create_nnx_sharded_model( - self.model, self._create_model_fn, mesh=self._mesh, named_sharding=full_sharding + self.model, self._create_model_fn, mesh=self._mesh, named_sharding=full_sharding # pyrefly: ignore[bad-argument-type] ) graphdef, _, _, rest_state = nnx.split(concrete_model, nnx.Param, nnx.Cache, ...) self.graphdef = graphdef @@ -435,7 +435,7 @@ def _load_params_nnx(self, params, rng): # PREFILL/AR attention ops have different cache variable shapes, and a # mismatch trips the `assert prefill_kv_cache` check inside attention_op. with nn_partitioning.axis_rules(self.config.logical_axis_rules): - concrete_model = self._create_model_fn() + concrete_model = self._create_model_fn() # pyrefly: ignore[not-callable] graphdef, _, _, rest_state = nnx.split(concrete_model, nnx.Param, nnx.Cache, ...) # Overlay loaded non-Param/non-Cache leaves (e.g. AQT qrhs.frozen) onto # the PREFILL-mode rest_state. The PREFILL concrete_model already has @@ -443,8 +443,8 @@ def _load_params_nnx(self, params, rng): # values. Anything only in `loaded_rest_state` (e.g. AR-only RNG slots) # is ignored. We keep PREFILL rest_state as the base so RNG variables # match the PREFILL graphdef's expectations. - loaded_rest_dict = loaded_rest_state.to_pure_dict() - rest_dict = rest_state.to_pure_dict() + loaded_rest_dict = loaded_rest_state.to_pure_dict() # pyrefly: ignore[missing-attribute] + rest_dict = rest_state.to_pure_dict() # pyrefly: ignore[missing-attribute] def _overlay(dst, src): if isinstance(dst, dict) and isinstance(src, dict): @@ -768,7 +768,7 @@ def _prefill_jit( one_d_output = ones_to_keep * DECODING_ACTIVE_SEQUENCE_INDICATOR sequence_indicator = jnp.expand_dims(one_d_output, 0) - rng, new_rng = jax.random.split(rng) + rng, new_rng = jax.random.split(rng) # pyrefly: ignore[bad-argument-type] if self.config.pure_nnx: # Prefill always operates on batch=1 (one padded prompt at a time). nnx_cache = ( @@ -1026,7 +1026,7 @@ def _prefill_multisampling_jit( one_d_output = ones_to_keep * DECODING_ACTIVE_SEQUENCE_INDICATOR sequence_indicator = jnp.expand_dims(one_d_output, 0) - rng, new_rng = jax.random.split(rng) + rng, new_rng = jax.random.split(rng) # pyrefly: ignore[bad-argument-type] if self.config.pure_nnx: # Prefill is batch=1 (one prompt); multi-sampling only draws several first # tokens from the shared logits below. Mirror the _prefill_jit NNX branch. diff --git a/src/maxtext/input_pipeline/grain_data_processing.py b/src/maxtext/input_pipeline/grain_data_processing.py index 9141f3bb1c..a1aa47705a 100644 --- a/src/maxtext/input_pipeline/grain_data_processing.py +++ b/src/maxtext/input_pipeline/grain_data_processing.py @@ -204,9 +204,9 @@ def create_dataset_from_pattern(pattern): dataset = dataset.repeat(num_epoch) dataset = dataset[file_slice] if data_file_type == "tfrecord": - dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset) + dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset) # pyrefly: ignore[missing-attribute] else: - dataset = dataset.map(grain.experimental.ParquetIterDataset) + dataset = dataset.map(grain.experimental.ParquetIterDataset) # pyrefly: ignore[missing-attribute] cycle_length = min(files_per_host, grain_num_threads) dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=cycle_length) if row_shard is not None: diff --git a/src/maxtext/input_pipeline/grain_tokenizer.py b/src/maxtext/input_pipeline/grain_tokenizer.py index 0ca4d222db..06556f4027 100644 --- a/src/maxtext/input_pipeline/grain_tokenizer.py +++ b/src/maxtext/input_pipeline/grain_tokenizer.py @@ -71,7 +71,7 @@ class TokenizeAndTrim(TokenizerTransformBase, grain.MapTransform): def map(self, element: dict[str, Any]) -> dict[str, Any]: """Maps to each element.""" - for feature_name, max_length in zip(self.feature_names, self.sequence_length, strict=True): + for feature_name, max_length in zip(self.feature_names, self.sequence_length, strict=True): # pyrefly: ignore[bad-argument-type] text = element[feature_name] token_ids = self._encode(text)[:max_length] element[feature_name] = np.asarray(token_ids, dtype=np.int32) @@ -88,9 +88,9 @@ def __post_init__(self): super().__post_init__() # TokenizeAndChunk only supports single feature for chunking assert len(self.feature_names) == 1, "TokenizeAndChunk only supports single feature name" - assert len(self.sequence_length) == 1, "TokenizeAndChunk only supports single sequence length" + assert len(self.sequence_length) == 1, "TokenizeAndChunk only supports single sequence length" # pyrefly: ignore[bad-argument-type] self.feature_name = self.feature_names[0] # For backward compatibility - self.sequence_length = self.sequence_length[0] # Convert back to int for chunking + self.sequence_length = self.sequence_length[0] # Convert back to int for chunking # pyrefly: ignore[bad-index] def flat_map(self, element: dict[str, Any]) -> list[dict[str, Any]]: """Tokenize and chunk text into multiple examples of sequence length.""" @@ -103,8 +103,8 @@ def flat_map(self, element: dict[str, Any]) -> list[dict[str, Any]]: return [] output_elements = [] - for start_idx in range(0, len(token_ids), chunk_size): - chunk = np.asarray(token_ids[start_idx : start_idx + chunk_size], dtype=np.int32) + for start_idx in range(0, len(token_ids), chunk_size): # pyrefly: ignore[bad-argument-type] + chunk = np.asarray(token_ids[start_idx : start_idx + chunk_size], dtype=np.int32) # pyrefly: ignore[unsupported-operation] new_element = {self.feature_name: chunk} output_elements.append(new_element) diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index 8ba57363fb..0b7751b522 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/src/maxtext/input_pipeline/input_pipeline_utils.py @@ -172,7 +172,7 @@ def is_conversational(features, data_columns): for column in data_columns: messages = features[column] if isinstance(messages, datasets.Sequence): - if isinstance(messages.feature, dict) and "role" in messages.feature and "content" in messages.feature: + if isinstance(messages.feature, dict) and "role" in messages.feature and "content" in messages.feature: # pyrefly: ignore[missing-attribute] return True return False @@ -778,7 +778,7 @@ def _pad_image_and_mask(self, preprocessed_image: mm_utils.PreprocessorOutput) - if preprocessed_image.pixel_values is None: raise ValueError("Input preprocessed_image must have pixel_values to pad images.") - if self.config.model_name and self.config.model_name.startswith("qwen3-omni"): + if self.config.model_name and self.config.model_name.startswith("qwen3-omni"): # pyrefly: ignore[missing-attribute] return preprocessed_image # Determine the maximum number of images/masks allowed. @@ -840,21 +840,21 @@ def map( if isinstance(element[data_column], mm_utils.PreprocessorOutput): raise TypeError("Only 'images' column can be of type PreprocessorOutput.") - element[f"{data_column}_segmentation"] = element[data_column] != self.pad_id - element[f"{data_column}_segmentation"] = element[f"{data_column}_segmentation"].astype(np.int32) - element[f"{data_column}_position"] = np.arange(element[data_column].shape[0], dtype=np.int32) + element[f"{data_column}_segmentation"] = element[data_column] != self.pad_id # pyrefly: ignore[unsupported-operation] + element[f"{data_column}_segmentation"] = element[f"{data_column}_segmentation"].astype(np.int32) # pyrefly: ignore[missing-attribute] + element[f"{data_column}_position"] = np.arange(element[data_column].shape[0], dtype=np.int32) # pyrefly: ignore[missing-attribute] if self.add_true_length: - element[f"{data_column}_true_length"] = np.array([element[data_column].shape[0]], dtype=np.int32) + element[f"{data_column}_true_length"] = np.array([element[data_column].shape[0]], dtype=np.int32) # pyrefly: ignore[missing-attribute] for key, _ in element.items(): if key == "images": - if self.config.model_name is None: + if self.config.model_name is None: # pyrefly: ignore[missing-attribute] raise ValueError("model_name must be provided when padding images") - element["images"] = self._pad_image_and_mask(element["images"]) + element["images"] = self._pad_image_and_mask(element["images"]) # pyrefly: ignore[bad-argument-type] elif "true_length" not in key: - element[key] = self._pad_text(element[key], self.max_length, self.pad_id) + element[key] = self._pad_text(element[key], self.max_length, self.pad_id) # pyrefly: ignore[bad-argument-type] return element @@ -881,7 +881,7 @@ def map(self, element: dict[str, np.ndarray]) -> dict[str, np.ndarray]: raise TypeError(f"'images' must be of type PreprocessorOutput, but got {type(preprocessed_image)}") output = element.copy() - output["images"] = preprocessed_image.pixel_values + output["images"] = preprocessed_image.pixel_values # pyrefly: ignore[unsupported-operation] if preprocessed_image.pixel_mask is not None: output["image_masks"] = preprocessed_image.pixel_mask diff --git a/src/maxtext/input_pipeline/multihost_dataloading.py b/src/maxtext/input_pipeline/multihost_dataloading.py index 41f40a7608..8e1246a25b 100644 --- a/src/maxtext/input_pipeline/multihost_dataloading.py +++ b/src/maxtext/input_pipeline/multihost_dataloading.py @@ -239,11 +239,11 @@ def save_state(self, step_array): if jax.process_index() == 0: directory.mkdir(parents=True, exist_ok=True) filename = directory / "process_0.json" - filename.write_text(json.dumps(self.iterator.get_state(), indent=4)) + filename.write_text(json.dumps(self.iterator.get_state(), indent=4)) # pyrefly: ignore[missing-attribute] return step_array directory.mkdir(parents=True, exist_ok=True) filename = directory / f"process_{jax.process_index()}-of-{jax.process_count()}.json" - state = json.dumps(self.iterator.get_state(), indent=4) + state = json.dumps(self.iterator.get_state(), indent=4) # pyrefly: ignore[missing-attribute] filename.write_text(state) return step_array @@ -255,7 +255,7 @@ def restore_state(self, step_array): else: filename = directory / f"process_{jax.process_index()}-of-{jax.process_count()}.json" state = json.loads(filename.read_text()) - self.iterator.set_state(state) + self.iterator.set_state(state) # pyrefly: ignore[missing-attribute] return step_array @@ -273,11 +273,11 @@ def __init__(self, get_ds_fn, preprocessing_fn, global_mesh, global_shape, check # named "local_iterator" to match MultiHostDataLoadIterator's interface. remote_iterator_cls = colocated_python.colocated_python_class(RemoteIterator) self.local_iterator = remote_iterator_cls( - get_ds_fn, + get_ds_fn, # pyrefly: ignore[bad-argument-count] preprocessing_fn, global_shape, checkpoint_path, - elastic=elastic, + elastic=elastic, # pyrefly: ignore[unexpected-keyword] ) max_logging.log("RemoteIteratorWrapper initiated") @@ -285,10 +285,10 @@ def __iter__(self): return self def reset(self): - self.local_iterator.reset() + self.local_iterator.reset() # pyrefly: ignore[missing-attribute] def __next__(self): - out = self.local_iterator.get_next(self.dummy_array) + out = self.local_iterator.get_next(self.dummy_array) # pyrefly: ignore[missing-attribute] # use tree_map is out is a dict return jax.device_put(out, self.tpu_sharding) @@ -296,10 +296,10 @@ def save_state(self, step): replicated_cpu_sharding = NamedSharding(self.cpu_mesh, PartitionSpec()) step_array = jnp.array(step, dtype=jnp.int32) step_array = jax.device_put(step_array, replicated_cpu_sharding) - self.local_iterator.save_state(step_array) + self.local_iterator.save_state(step_array) # pyrefly: ignore[missing-attribute] def restore_state(self, step): replicated_cpu_sharding = NamedSharding(self.cpu_mesh, PartitionSpec()) step_array = jnp.array(step, dtype=jnp.int32) step_array = jax.device_put(step_array, replicated_cpu_sharding) - self.local_iterator.restore_state(step_array) + self.local_iterator.restore_state(step_array) # pyrefly: ignore[missing-attribute] diff --git a/src/maxtext/input_pipeline/olmo_data.py b/src/maxtext/input_pipeline/olmo_data.py index 82c258b78a..39d5cdd031 100644 --- a/src/maxtext/input_pipeline/olmo_data.py +++ b/src/maxtext/input_pipeline/olmo_data.py @@ -137,7 +137,7 @@ def global_to_local(index: OlmoNpyIndex, instance_id: int) -> Tuple[int, int]: if instance_id < 0 or instance_id >= index.total_instances: raise IndexError(f"instance_id {instance_id} out of range " f"[0, {index.total_instances})") starts = index._instance_offset_starts # type: ignore[attr-defined] # pylint: disable=protected-access - file_idx = bisect.bisect_right(starts, instance_id) - 1 + file_idx = bisect.bisect_right(starts, instance_id) - 1 # pyrefly: ignore[bad-argument-type] local_instance = instance_id - index.files[file_idx].instance_offset token_offset = local_instance * index.sequence_length return file_idx, token_offset diff --git a/src/maxtext/input_pipeline/packing/prefill_packing.py b/src/maxtext/input_pipeline/packing/prefill_packing.py index 8cc8816437..d3f74f5ab2 100644 --- a/src/maxtext/input_pipeline/packing/prefill_packing.py +++ b/src/maxtext/input_pipeline/packing/prefill_packing.py @@ -330,7 +330,7 @@ def zero_padded(arr: list[int], padding: int): prompt_logp_numpy = np.array(decode_state["prompt_logp"]) for i in range(bucket.count): if return_prompt_logp: - prompt_logp = prompt_logp_numpy[:, offsets[i] : offsets[i] + lengths[i]] + prompt_logp = prompt_logp_numpy[:, offsets[i] : offsets[i] + lengths[i]] # pyrefly: ignore[unsupported-operation] prefill_result.append(PrefillResult(first_tokens[i], bucket.slots[i], prompt_logp)) else: prefill_result.append(PrefillResult(first_tokens[i], bucket.slots[i], None)) diff --git a/src/maxtext/kernels/attention/splash_attention_kernel.py b/src/maxtext/kernels/attention/splash_attention_kernel.py index bd5099972f..0eb57087f6 100644 --- a/src/maxtext/kernels/attention/splash_attention_kernel.py +++ b/src/maxtext/kernels/attention/splash_attention_kernel.py @@ -114,7 +114,7 @@ def get_kernel_name( @overload -def _attention_reference( +def _attention_reference( # pyrefly: ignore[inconsistent-overload] mask: jax.Array, q: jax.Array, k: jax.Array, @@ -130,7 +130,7 @@ def _attention_reference( @overload -def _attention_reference( +def _attention_reference( # pyrefly: ignore[inconsistent-overload] mask: jax.Array, q: jax.Array, k: jax.Array, @@ -618,18 +618,18 @@ def _apply_mask_and_soft_cap( # need to keep into account the current shard along Q sequence. if k_in_lanes: - assert q_sequence_ref.shape == (bq, NUM_LANES) + assert q_sequence_ref.shape == (bq, NUM_LANES) # pyrefly: ignore[missing-attribute] k_sequence = k_offset + jax.lax.broadcasted_iota(jnp.int32, (bq, k_slice.size), 1) repeats, rem = divmod(k_slice.size, NUM_LANES) assert rem == 0 - q_sequence = jnp.tile(q_sequence_ref[...], (1, repeats)) # [bq, k_slice.size] + q_sequence = jnp.tile(q_sequence_ref[...], (1, repeats)) # [bq, k_slice.size] # pyrefly: ignore[unsupported-operation] else: - assert q_sequence_ref.shape == (NUM_SUBLANES, bq) + assert q_sequence_ref.shape == (NUM_SUBLANES, bq) # pyrefly: ignore[missing-attribute] k_sequence = k_offset + jax.lax.broadcasted_iota(jnp.int32, (k_slice.size, bq), 0) - q_sequence = q_sequence_ref[:1, :] # [1, bq] + q_sequence = q_sequence_ref[:1, :] # [1, bq] # pyrefly: ignore[unsupported-operation] q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq)) assert q_sequence.shape == k_sequence.shape @@ -773,7 +773,7 @@ def body(kv_compute_index, _): if rem != 0: raise NotImplementedError(f"{bkv_compute=} should be a multiple of {NUM_LANES}") - s_curr = jnp.exp(qk - jnp.tile(m_next, (1, bkv_repeats))) + s_curr = jnp.exp(qk - jnp.tile(m_next, (1, bkv_repeats))) # pyrefly: ignore[unsupported-operation] assert s_curr.shape == (bq, bkv_compute) l_curr = jax.lax.broadcast_in_dim(s_curr.sum(axis=-1), l_prev.shape, (0,)) @@ -815,7 +815,7 @@ def end(): @overload -def _splash_attention_forward( +def _splash_attention_forward( # pyrefly: ignore[inconsistent-overload] fwd_mask_info: mask_info_lib.MaskInfo, q: jax.Array, k: jax.Array, @@ -833,7 +833,7 @@ def _splash_attention_forward( @overload -def _splash_attention_forward( +def _splash_attention_forward( # pyrefly: ignore[inconsistent-overload] fwd_mask_info: mask_info_lib.MaskInfo, q: jax.Array, k: jax.Array, @@ -908,9 +908,9 @@ def _splash_attention_forward( if k.shape[:-1] != v.shape[:-1]: raise ValueError(f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " "leading dimensions.") - if bkv % bkv_compute: + if bkv % bkv_compute: # pyrefly: ignore[unsupported-operation] raise ValueError(f"{bkv=} must be a multiple of {bkv_compute=}.") - if bkv_compute % NUM_LANES: + if bkv_compute % NUM_LANES: # pyrefly: ignore[unsupported-operation] raise ValueError(f"{bkv_compute=} must be a multiple of {NUM_LANES}.") kv_seq_len = k.shape[kv_seq_len_dimension] @@ -988,7 +988,7 @@ def kv_segment_ids_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_r if fwd_mask_info.partial_mask_blocks is not None: in_specs.append(pl.BlockSpec((None, bq, bkv), mask_index_map)) else: - in_specs.append(None) + in_specs.append(None) # pyrefly: ignore[bad-argument-type] assert fwd_mask_info.partial_mask_blocks is None or fwd_mask_info.q_sequence is None @@ -997,7 +997,7 @@ def kv_segment_ids_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_r in_specs.append(pl.BlockSpec((bq, NUM_LANES), q_segment_ids_index_map)) else: q_sequence = None - in_specs.append(None) + in_specs.append(None) # pyrefly: ignore[bad-argument-type] num_scalar_prefetch = 3 @@ -1266,8 +1266,8 @@ def run(): q_sequence_ref, q_segment_ids_ref, kv_segment_ids_ref, - attn_logits_soft_cap=attn_logits_soft_cap, - k_slice=pl.ds(0, bkv), + attn_logits_soft_cap=attn_logits_soft_cap, # pyrefly: ignore[bad-argument-type] + k_slice=pl.ds(0, bkv), # pyrefly: ignore[bad-argument-type] # When the iteration space is shrunk (for local attention for example), # the kv_index program_id does not correspond to the actual coordinates # of the KV data. Make sure to use the 'unshrunk' index (coming from the @@ -1276,7 +1276,7 @@ def run(): bq=bq, mask_function=mask_function, ) - p = jnp.exp(qk - logsumexp) + p = jnp.exp(qk - logsumexp) # pyrefly: ignore[unsupported-operation] dp_dims = NT_DIM_NUMBERS if v_layout == HEAD_DIM_MINOR else NN_DIM_NUMBERS dp = lax.dot_general( do.astype(v.dtype), @@ -1426,11 +1426,11 @@ def logsumexp_index_map(h, i, *_): logsumexp = jnp.expand_dims(logsumexp, axis=-2) logsumexp_spec = pl.BlockSpec((None, 1, bq), logsumexp_index_map) - assert logsumexp.ndim == len(logsumexp_spec.block_shape) + assert logsumexp.ndim == len(logsumexp_spec.block_shape) # pyrefly: ignore[bad-argument-type] di = jnp.expand_dims(di, axis=-2) di_spec = pl.BlockSpec((None, 1, bq), logsumexp_index_map) - assert di.ndim == len(di_spec.block_shape) + assert di.ndim == len(di_spec.block_shape) # pyrefly: ignore[bad-argument-type] in_specs = [ q_spec, @@ -1633,8 +1633,8 @@ def _load_kv(ref, layout): q_sequence_ref, q_segment_ids_ref, kv_segment_ids_ref, - attn_logits_soft_cap=attn_logits_soft_cap, - k_slice=slice_k, + attn_logits_soft_cap=attn_logits_soft_cap, # pyrefly: ignore[bad-argument-type] + k_slice=slice_k, # pyrefly: ignore[bad-argument-type] k_offset=kv_index * bkv + i * bkv_compute, bq=bq, k_in_lanes=False, @@ -1694,7 +1694,7 @@ def run(): if is_mqa: should_write = jnp.logical_and(should_write, q_head_index == num_q_heads - 1) elif num_kv_heads < num_q_heads: - should_write = jnp.logical_and(should_write, q_head_index_per_kv_head == q_heads_per_kv_heads - 1) + should_write = jnp.logical_and(should_write, q_head_index_per_kv_head == q_heads_per_kv_heads - 1) # pyrefly: ignore[unsupported-operation] @pl.when(should_write) def end(): @@ -1943,12 +1943,12 @@ def logsumexp_index_map( logsumexp_shape = (num_q_heads, NUM_SUBLANES, q_seq_len) logsumexp = jnp.broadcast_to(jnp.expand_dims(logsumexp, -2), logsumexp_shape) logsumexp_spec = pl.BlockSpec((None, NUM_SUBLANES, bq), logsumexp_index_map) - assert logsumexp.ndim == len(logsumexp_spec.block_shape) + assert logsumexp.ndim == len(logsumexp_spec.block_shape) # pyrefly: ignore[bad-argument-type] # TODO(apaszke): Remove the sublane expansion once Mosaic has all retilings di = jnp.broadcast_to(jnp.expand_dims(di, -2), logsumexp_shape) di_spec = pl.BlockSpec((None, NUM_SUBLANES, bq), logsumexp_index_map) - assert di.ndim == len(di_spec.block_shape) + assert di.ndim == len(di_spec.block_shape) # pyrefly: ignore[bad-argument-type] in_specs = [ q_spec, @@ -2117,11 +2117,11 @@ def _splash_attention_bwd( logsumexp, do, di, - bq=bq_dkv, - bkv=bkv_dkv_memory, - bkv_compute=bkv_dkv_compute, + bq=bq_dkv, # pyrefly: ignore[bad-argument-type] + bkv=bkv_dkv_memory, # pyrefly: ignore[bad-argument-type] + bkv_compute=bkv_dkv_compute, # pyrefly: ignore[bad-argument-type] is_mqa=is_mqa, - mask_info=dkv_mask_info, + mask_info=dkv_mask_info, # pyrefly: ignore[bad-argument-type] mask_value=mask_value, attn_logits_soft_cap=attn_logits_soft_cap, use_fused_bwd_kernel=use_fused_bwd_kernel, @@ -2141,10 +2141,10 @@ def _splash_attention_bwd( logsumexp, do, di, - bq=bq_dq, - bkv=bkv_dq, + bq=bq_dq, # pyrefly: ignore[bad-argument-type] + bkv=bkv_dq, # pyrefly: ignore[bad-argument-type] is_mqa=is_mqa, - mask_info=dq_mask_info, + mask_info=dq_mask_info, # pyrefly: ignore[bad-argument-type] mask_value=mask_value, attn_logits_soft_cap=attn_logits_soft_cap, q_layout=block_sizes.q_layout, @@ -2348,7 +2348,7 @@ def _splash_attention_manual_bwd( save_residuals=save_residuals, mask_value=mask_value, is_mqa=is_mqa, - block_sizes=block_sizes, + block_sizes=block_sizes, # pyrefly: ignore[bad-argument-type] residual_checkpoint_name=residual_checkpoint_name, mask_function=mask_function, attn_logits_soft_cap=attn_logits_soft_cap, @@ -2422,11 +2422,11 @@ def manual_sharding_spec(self, sharding: jax.sharding.NamedSharding): # Shard q_sequence over the sequence dimension only. q_sequence_spec = jax.sharding.PartitionSpec(spec[1]) mask_info_specs = mask_info_lib.MaskInfo( # pytype: disable=wrong-arg-types - data_next=spec if self.fwd_mask_info.data_next is not None else None, - mask_next=spec if self.fwd_mask_info.mask_next is not None else None, - block_mask=spec if self.fwd_mask_info.block_mask is not None else None, - partial_mask_blocks=partial_mask_blocks_spec if self.fwd_mask_info.partial_mask_blocks is not None else None, - q_sequence=q_sequence_spec if self.fwd_mask_info.q_sequence is not None else None, + data_next=spec if self.fwd_mask_info.data_next is not None else None, # pyrefly: ignore[bad-argument-type] + mask_next=spec if self.fwd_mask_info.mask_next is not None else None, # pyrefly: ignore[bad-argument-type] + block_mask=spec if self.fwd_mask_info.block_mask is not None else None, # pyrefly: ignore[bad-argument-type] + partial_mask_blocks=partial_mask_blocks_spec if self.fwd_mask_info.partial_mask_blocks is not None else None, # pyrefly: ignore[bad-argument-type] + q_sequence=q_sequence_spec if self.fwd_mask_info.q_sequence is not None else None, # pyrefly: ignore[bad-argument-type] ) return SplashAttentionKernel( mask_info_specs, diff --git a/src/maxtext/kernels/gather_reduce_sc.py b/src/maxtext/kernels/gather_reduce_sc.py index 5b3b8e7597..792f799e33 100644 --- a/src/maxtext/kernels/gather_reduce_sc.py +++ b/src/maxtext/kernels/gather_reduce_sc.py @@ -473,7 +473,7 @@ def perform_add( is_odd = arith.cmpi( arith.CmpIPredicate.eq, - parity, + parity, # pyrefly: ignore[bad-argument-type] const_lut(0), ) raw_weights_f32 = arith.select(is_odd, weights_odds, weights_evens) @@ -533,8 +533,8 @@ def get_row_val(row_idx): _BF16[2, 16], scratch_local, [ - arith.addi(const_lut(row_idx_l), row_add), - arith.muli(col_pos, const_lut(2)), # NOTYPO + arith.addi(const_lut(row_idx_l), row_add), # pyrefly: ignore[bad-argument-type] + arith.muli(col_pos, const_lut(2)), # NOTYPO # pyrefly: ignore[bad-argument-type] ], enable_all_sublanes_mask, ) @@ -567,20 +567,20 @@ def get_row_val(row_idx): row0 = get_row_val(0) if weights_local is not None: - row0 = arith.mulf(row0, weights_vecs[0]) + row0 = arith.mulf(row0, weights_vecs[0]) # pyrefly: ignore[unsupported-operation] if topk_wgt_zero_nan: row0 = arith.select( - arith.cmpf(arith.CmpFPredicate.OEQ, weights_vecs[0], zero_vec_f32), + arith.cmpf(arith.CmpFPredicate.OEQ, weights_vecs[0], zero_vec_f32), # pyrefly: ignore[unbound-name, unsupported-operation] zero_vec_f32, row0, ) row8 = get_row_val(8) if weights_local is not None: - row8 = arith.mulf(row8, weights_vecs[8]) + row8 = arith.mulf(row8, weights_vecs[8]) # pyrefly: ignore[unsupported-operation] if topk_wgt_zero_nan: row8 = arith.select( - arith.cmpf(arith.CmpFPredicate.OEQ, weights_vecs[8], zero_vec_f32), + arith.cmpf(arith.CmpFPredicate.OEQ, weights_vecs[8], zero_vec_f32), # pyrefly: ignore[unsupported-operation] zero_vec_f32, row8, ) @@ -588,12 +588,12 @@ def get_row_val(row_idx): for sum_idx in range(7): tmp_row0 = get_row_val(sum_idx + 1) if weights_local is not None: - tmp_row0 = arith.mulf(tmp_row0, weights_vecs[sum_idx + 1]) + tmp_row0 = arith.mulf(tmp_row0, weights_vecs[sum_idx + 1]) # pyrefly: ignore[unsupported-operation] if topk_wgt_zero_nan: tmp_row0 = arith.select( arith.cmpf( arith.CmpFPredicate.OEQ, - weights_vecs[sum_idx + 1], + weights_vecs[sum_idx + 1], # pyrefly: ignore[unsupported-operation] zero_vec_f32, ), zero_vec_f32, @@ -604,12 +604,12 @@ def get_row_val(row_idx): tmp_row8 = get_row_val(8 + sum_idx + 1) if weights_local is not None: - tmp_row8 = arith.mulf(tmp_row8, weights_vecs[8 + sum_idx + 1]) + tmp_row8 = arith.mulf(tmp_row8, weights_vecs[8 + sum_idx + 1]) # pyrefly: ignore[unsupported-operation] if topk_wgt_zero_nan: tmp_row8 = arith.select( arith.cmpf( arith.CmpFPredicate.OEQ, - weights_vecs[8 + sum_idx + 1], + weights_vecs[8 + sum_idx + 1], # pyrefly: ignore[unsupported-operation] zero_vec_f32, ), zero_vec_f32, @@ -743,7 +743,7 @@ def fill_out_offset_tile(offset_tile_out_local, col_pos, row_pos=None): ) if row_pos is not None: - local_vec = arith.addi(vec, row_vec_offset) + local_vec = arith.addi(vec, row_vec_offset) # pyrefly: ignore[unbound-name] else: local_vec = vec @@ -1154,7 +1154,7 @@ def fill_out_offset_tile(offset_tile_out_local, col_pos, row_pos=None): lower_bound=const_lut(row_chunk_size * 1), upper_bound=const_lut(row_chunk_size * ((idx.shape[0] // num_sc // row_chunk_size) - 1)), step=const_lut(row_chunk_size * 2), - iter_args=[ + iter_args=[ # pyrefly: ignore[bad-argument-type] scratch_0, scratch_out_0, sflag_0, diff --git a/src/maxtext/kernels/megablox/backend.py b/src/maxtext/kernels/megablox/backend.py index e736bdb000..de3fab8b92 100644 --- a/src/maxtext/kernels/megablox/backend.py +++ b/src/maxtext/kernels/megablox/backend.py @@ -345,7 +345,7 @@ def gmm( group_offset = group_offset[None] num_current_groups = rhs.shape[0] num_total_groups = group_sizes.shape[0] - group_sizes = _validate_args(lhs=lhs, rhs=rhs, group_sizes=group_sizes) + group_sizes = _validate_args(lhs=lhs, rhs=rhs, group_sizes=group_sizes) # pyrefly: ignore[bad-argument-type] # Gather shape information. m, k, n = (lhs.shape[0], lhs.shape[1], rhs.shape[2]) @@ -607,7 +607,7 @@ def tgmm( group_offset = jnp.array([0], dtype=jnp.int32) else: group_offset = group_offset[None] - group_sizes = _validate_args(lhs=lhs, rhs=rhs, group_sizes=group_sizes, expected_rhs_dims=2) + group_sizes = _validate_args(lhs=lhs, rhs=rhs, group_sizes=group_sizes, expected_rhs_dims=2) # pyrefly: ignore[bad-argument-type] # Gather shape information. k, m, n = (lhs.shape[0], lhs.shape[1], rhs.shape[1]) diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index b7487c7368..06f9aa9eac 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -63,8 +63,8 @@ def gmm( existing_out: jnp.ndarray | None = None, transpose_rhs: bool = False, interpret: bool = False, - lhs_quantize_dtype: Literal[jnp.int4, jnp.int8] | None = None, - rhs_quantize_dtype: Literal[jnp.int4, jnp.int8] | None = None, + lhs_quantize_dtype: Literal[jnp.int4, jnp.int8] | None = None, # pyrefly: ignore[invalid-literal] + rhs_quantize_dtype: Literal[jnp.int4, jnp.int8] | None = None, # pyrefly: ignore[invalid-literal] use_qwix_quantization: bool = False, use_tokamax_backend: bool = False, weight_gather_axes: List[Tuple[str, int]] | None = None, @@ -154,15 +154,15 @@ def _gmm_fwd( """Forward function for GMM VJP.""" if quantization_rule: if quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray): - lhs = qpl.quantize( + lhs = qpl.quantize( # pyrefly: ignore[bad-assignment] lhs, quantization_rule.act_qtype, channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0], - calibration_method=quantization_rule.act_calibration_method, + calibration_method=quantization_rule.act_calibration_method, # pyrefly: ignore[bad-argument-type] ) if quantization_rule.weight_qtype and not isinstance(rhs, qpl.QArray): if not use_manual_quantization: - rhs = qpl.quantize( + rhs = qpl.quantize( # pyrefly: ignore[bad-assignment] rhs, quantization_rule.weight_qtype, # If only considering the fwd pass, we could also enable channelwise @@ -172,7 +172,7 @@ def _gmm_fwd( calibration_method=quantization_rule.weight_calibration_method, ) else: - rhs = quantizations.manual_quantize( + rhs = quantizations.manual_quantize( # pyrefly: ignore[bad-assignment] rhs, quantization_rule.weight_qtype, calibration_method=quantization_rule.weight_calibration_method, diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index 8240e1bf14..c1ebeda0e3 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py @@ -210,7 +210,7 @@ class GmmConfigs: @property def num_quant_blocks_per_tile_k(self) -> int: - return pl.cdiv(self.tiles.tile_k, self.rhs_cfgs.quant_block_size) + return pl.cdiv(self.tiles.tile_k, self.rhs_cfgs.quant_block_size) # pyrefly: ignore[no-matching-overload] @property def out_size_n(self) -> int: @@ -253,7 +253,7 @@ def rhs_scale_index_map(self, n_id: jax.Array, gm_id: jax.Array, k_id: jax.Array # Simply multiplying k_id by num_quant_blocks_per_tile_k will not work # since a single quant block could be shared along multiple k tile. k_row = k_id * self.cfgs.tiles.tile_k - b_row = k_row // self.cfgs.rhs_cfgs.quant_block_size + b_row = k_row // self.cfgs.rhs_cfgs.quant_block_size # pyrefly: ignore[unsupported-operation] b_tile_id = b_row // self.cfgs.num_quant_blocks_per_tile_k return (group_id, b_tile_id, 0, n_id) @@ -420,8 +420,8 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): acc_n = jnp.zeros((cfgs.tiles.tile_m, col_size), dtype=acc_ref.dtype) for b_id in range(cfgs.num_quant_blocks_per_tile_k): - start_k = b_id * rhs_qbs - end_k = start_k + rhs_qbs + start_k = b_id * rhs_qbs # pyrefly: ignore[unsupported-operation] + end_k = start_k + rhs_qbs # pyrefly: ignore[unsupported-operation] block_acc = jnp.matmul( tiled_lhs[:, start_k:end_k], @@ -458,8 +458,8 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): col_size = end_n - start_n acc_n = jnp.zeros((cfgs.tiles.tile_m, col_size), dtype=acc_ref.dtype) - for start_k in range(0, cfgs.tiles.tile_k, q_block_size): - end_k = min(cfgs.tiles.tile_k, start_k + q_block_size) + for start_k in range(0, cfgs.tiles.tile_k, q_block_size): # pyrefly: ignore[bad-argument-type] + end_k = min(cfgs.tiles.tile_k, start_k + q_block_size) # pyrefly: ignore[unsupported-operation] block_lhs = tiled_lhs[:, start_k:end_k] block_rhs = tiled_rhs[start_k:end_k, start_n:end_n] @@ -493,7 +493,7 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): # Apply rhs subchannel scale per quant block. if cfgs.rhs_cfgs.should_dequantize_after_matmul: - b_id = start_k // cfgs.rhs_cfgs.quant_block_size + b_id = start_k // cfgs.rhs_cfgs.quant_block_size # pyrefly: ignore[unsupported-operation] rhs_scale_slice = tiled_rhs_ref.get_scale() block_acc *= rhs_scale_slice[b_id, :, start_n:end_n].astype(acc_ref.dtype) @@ -510,7 +510,7 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): tiled_rhs_bias = tiled_rhs_ref.get_bias() acc += tiled_rhs_bias.astype(acc.dtype) if cfgs.has_partial_sum: - ps_tile = tiled_ps_ref[...].reshape(acc.shape) + ps_tile = tiled_ps_ref[...].reshape(acc.shape) # pyrefly: ignore[unsupported-operation] acc += ps_tile.astype(acc.dtype) acc = apply_act_fn(acc, cfgs.fuse_act) @@ -829,8 +829,8 @@ def kernel_main( if cfgs.zero_init: zero_size = zero_out_start( out_ref, - zero_ref, - semaphore_ref, + zero_ref, # pyrefly: ignore[bad-argument-type] + semaphore_ref, # pyrefly: ignore[bad-argument-type] metadata_ref, num_gm, dims=cfgs.dims, @@ -840,7 +840,7 @@ def kernel_main( if cfgs.fuse_act is not None: rhs_up_ref = jax.tree.map(lambda x: x.at[..., cfgs.out_size_n :], rhs_ref) - rhs_ref = FusedWeightsRef(gate=rhs_ref, up=rhs_up_ref) + rhs_ref = FusedWeightsRef(gate=rhs_ref, up=rhs_up_ref) # pyrefly: ignore[bad-assignment] rhs_spec = FusedWeightsRef( gate=rhs_spec, @@ -866,7 +866,7 @@ def kernel_main( pipeline_fn(lhs_in, rhs_ref, ps_in, out_in, scratches=scratches) if cfgs.zero_init: - zero_out_end(out_ref, semaphore_ref, zero_size, dims=cfgs.dims) + zero_out_end(out_ref, semaphore_ref, zero_size, dims=cfgs.dims) # pyrefly: ignore[bad-argument-type, unbound-name] def calculate_tiling( @@ -909,7 +909,7 @@ def calculate_tiling( tile_n_limit //= fuse_act_factor def _is_tile_k_quant_block_compatible(tk: int) -> bool: - if tk % rhs_cfgs.quant_block_size != 0 and rhs_cfgs.quant_block_size % tk != 0: + if tk % rhs_cfgs.quant_block_size != 0 and rhs_cfgs.quant_block_size % tk != 0: # pyrefly: ignore[unsupported-operation] return False return True @@ -1046,7 +1046,7 @@ def get_cost_estimate(cfgs: GmmConfigs): rhs_size = dims.size_group * dims.size_k * dims.size_n rhs_bytes = rhs_size * rhs_bits // 8 if cfgs.rhs_cfgs.has_scale: - num_quant_blocks = pl.cdiv(dims.size_k, cfgs.rhs_cfgs.quant_block_size) + num_quant_blocks = pl.cdiv(dims.size_k, cfgs.rhs_cfgs.quant_block_size) # pyrefly: ignore[no-matching-overload] rhs_bytes += dims.size_group * num_quant_blocks * dims.size_n * fp32_bytes if cfgs.rhs_cfgs.has_bias: rhs_bytes += dims.size_group * dims.size_n * fp32_bytes @@ -1113,14 +1113,14 @@ def make_gmm_configs( lhs_q_dtype = None if maybe_quantize_lhs and rhs_cfgs.should_dequantize_after_matmul: # Choose lhs quantization dtype based on TPU hardware support. - is_rhs_float = jnp.issubdtype(rhs_quant_dtype, jnp.floating) + is_rhs_float = jnp.issubdtype(rhs_quant_dtype, jnp.floating) # pyrefly: ignore[bad-argument-type] tpu_info = pltpu.get_tpu_info() # Check if there is hardware compute support for rhs dtype group. if tpu_info.fp8_ops_per_second > 0: # Special handling for 4-bit integer rhs as it can be converted to fp8 # without a numeric issues. Note that this is not the case for 4-bit # floating rhs as conversion to int8 will cause numeric issues. - is_rhs_4bits = jax.dtypes.itemsize_bits(rhs_quant_dtype) == 4 + is_rhs_4bits = jax.dtypes.itemsize_bits(rhs_quant_dtype) == 4 # pyrefly: ignore[bad-argument-type] if is_rhs_float or is_rhs_4bits: lhs_q_dtype = jnp.float8_e4m3fn.dtype if tpu_info.int8_ops_per_second > 0: @@ -1199,7 +1199,7 @@ def gmm_v2( partial_sum: jax.Array | None = None, # [size_m, size_n] group_offset: jax.Array | None = None, # int32[1] *, - tile_info: TileSizes | TileFn = calculate_tiling, + tile_info: TileSizes | TileFn = calculate_tiling, # pyrefly: ignore[bad-function-definition] vmem_limit_bytes: int | None = None, precision: jax.lax.Precision = jax.lax.Precision.DEFAULT, preferred_element_type: jnp.dtype | None = None, @@ -1348,7 +1348,7 @@ def gmm_v2( num_scalar_prefetch=2, in_specs=in_specs, out_specs=pl.BlockSpec(memory_space=pltpu.HBM), - scratch_shapes=scratch_shapes, + scratch_shapes=scratch_shapes, # pyrefly: ignore[bad-argument-type] ), compiler_params=pltpu.CompilerParams( vmem_limit_bytes=vmem_limit_bytes, diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py index 01b17bd247..26fcdef05b 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py @@ -252,7 +252,7 @@ def make_tgmm_configs( dims, lhs_cfgs, rhs_cfgs, - vmem_limit_bytes, + vmem_limit_bytes, # pyrefly: ignore[bad-argument-type] out_dtype, acc_dtype, target_zero_ref_bytes, @@ -340,7 +340,7 @@ def _matmul(is_new_group: bool, is_group_changing: bool): if is_group_changing: if cfgs.rhs_cfgs.has_scale: - scale_slice = tiled_rhs_scale_ref[0] + scale_slice = tiled_rhs_scale_ref[0] # pyrefly: ignore[unsupported-operation] acc *= scale_slice tiled_out_ref[...] = acc.astype(tiled_out_ref.dtype) else: @@ -682,12 +682,12 @@ def tgmm_v2( cfgs = make_tgmm_configs( lhs, rhs, - rhs_scale, + rhs_scale, # pyrefly: ignore[bad-argument-type] group_sizes, num_actual_groups, tile_info=tile_info, vmem_limit_bytes=vmem_limit_bytes, - out_dtype=preferred_element_type, + out_dtype=preferred_element_type, # pyrefly: ignore[bad-argument-type] acc_dtype=acc_dtype, target_zero_ref_bytes=target_zero_ref_bytes, ) @@ -729,7 +729,7 @@ def tgmm_v2( pad_n = aligned_n - dims.size_n if pad_n > 0: rhs_scale = jnp.pad(rhs_scale, ((0, 0), (0, 0), (0, pad_n))) - rhs = OperandRef(value=rhs, scale=rhs_scale) + rhs = OperandRef(value=rhs, scale=rhs_scale) # pyrefly: ignore[bad-assignment] hbm_spec = pl.BlockSpec(memory_space=pltpu.HBM) in_specs = [ hbm_spec, # lhs @@ -745,7 +745,7 @@ def tgmm_v2( num_scalar_prefetch=2, in_specs=in_specs, out_specs=pl.BlockSpec(memory_space=pltpu.HBM), - scratch_shapes=scratch_shapes, + scratch_shapes=scratch_shapes, # pyrefly: ignore[bad-argument-type] ), compiler_params=pltpu.CompilerParams( vmem_limit_bytes=vmem_limit_bytes, diff --git a/src/maxtext/kernels/ragged/ragged_gather.py b/src/maxtext/kernels/ragged/ragged_gather.py index 14aa989c88..fc9aba84b8 100644 --- a/src/maxtext/kernels/ragged/ragged_gather.py +++ b/src/maxtext/kernels/ragged/ragged_gather.py @@ -140,8 +140,8 @@ def _(): # but for other bitwidths, this is not possible. Therefore, we bitcast data # into 32-bits first to fetch packing number of rows per dma and later # perform bitwise unpacking / packing to obtain desired results. - in_32b_hbm_ref = in_hbm_ref.bitcast(jnp.uint32) - out_32b_hbm_ref = out_hbm_ref.bitcast(jnp.uint32) + in_32b_hbm_ref = in_hbm_ref.bitcast(jnp.uint32) # pyrefly: ignore[missing-attribute] + out_32b_hbm_ref = out_hbm_ref.bitcast(jnp.uint32) # pyrefly: ignore[missing-attribute] for col_vmem_start in range(0, col_size, num_lanes): col_hbm_start = col_tile_start + col_vmem_start @@ -333,7 +333,7 @@ def _fallback_implementation( """Fallback to (non-ragged) JAX implementation for ragged gather.""" out = x[indices] if has_weights: - out = out * weights[:, None] + out = out * weights[:, None] # pyrefly: ignore[unsupported-operation] return out @@ -433,7 +433,7 @@ def ragged_gather( indices = jnp.pad(indices, ((0, out_pad_size))) if has_weights: - weights = jnp.pad(weights, ((0, out_pad_size)), constant_values=1.0) + weights = jnp.pad(weights, ((0, out_pad_size)), constant_values=1.0) # pyrefly: ignore[bad-argument-type] else: # Provide a dummy weights array; the kernel won't use it. weights = jnp.ones((out_size + out_pad_size,), dtype=jnp.float32) @@ -454,7 +454,7 @@ def ragged_gather( has_weights=has_weights, ), compiler_params=pltpu.CompilerParams( # pytype: disable=wrong-keyword-args - **_COMPILER_PARAMS, + **_COMPILER_PARAMS, # pyrefly: ignore[bad-argument-type] ), cost_estimate=get_cost_estimate( out_size=out_size + out_pad_size, @@ -466,7 +466,7 @@ def ragged_gather( ), mesh=vector_mesh, name="sc_ragged_gather", - **{ + **{ # pyrefly: ignore[bad-argument-type] _OUT_KW: jax.ShapeDtypeStruct((out_size + out_pad_size, aligned_hidden_size), dtype), _SCRATCH_KW: [ pltpu.VMEM((num_simd_lanes,), jnp.int32), diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index d5bb3314b6..d9385ce9d7 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -228,8 +228,8 @@ def row_loop(row_block_id): input_dtype_bits = jax.dtypes.itemsize_bits(in_dtype) input_packing = 32 // input_dtype_bits - in_32b_hbm_ref = in_hbm_ref.bitcast(jnp.uint32) - out_32b_hbm_ref = out_hbm_ref.bitcast(jnp.uint32) + in_32b_hbm_ref = in_hbm_ref.bitcast(jnp.uint32) # pyrefly: ignore[missing-attribute] + out_32b_hbm_ref = out_hbm_ref.bitcast(jnp.uint32) # pyrefly: ignore[missing-attribute] for col_vmem_start in range(0, col_size, num_lanes): col_hbm_start = col_start + col_vmem_start @@ -559,7 +559,7 @@ def ragged_gather_reduce( num_column_partitions=num_column_partitions, ), compiler_params=pltpu.CompilerParams( # pytype: disable=wrong-keyword-args - **_COMPILER_PARAMS, + **_COMPILER_PARAMS, # pyrefly: ignore[bad-argument-type] ), cost_estimate=get_cost_estimate( padded_input_size=padded_input_size, @@ -571,7 +571,7 @@ def ragged_gather_reduce( ), mesh=vector_mesh, name="sc_ragged_gather_reduce", - **{ + **{ # pyrefly: ignore[bad-argument-type] _OUT_KW: jax.ShapeDtypeStruct( (padded_input_size // reduce_group_size, aligned_hidden_size), jnp.float32, diff --git a/src/maxtext/layers/nnx_wrappers.py b/src/maxtext/layers/nnx_wrappers.py index a1d4ef64fe..2f5756023a 100644 --- a/src/maxtext/layers/nnx_wrappers.py +++ b/src/maxtext/layers/nnx_wrappers.py @@ -237,7 +237,7 @@ def __getattr__(self, name: str): maybe_method = getattr(type(self.to_nnx__module), name, None) if callable(maybe_method): method = partial(self.__call__, method=maybe_method) - method.__self__ = self + method.__self__ = self # pyrefly: ignore[missing-attribute] return method return super().__getattribute__(name) @@ -361,14 +361,14 @@ def wrapped(*args, **kwargs): if not linen.module._context.module_stack: # pylint: disable=W0212 return call_fn(*args, **kwargs) nn_module = linen.module._context.module_stack[-1] # pylint: disable=W0212 - old_path = nn_module.path + old_path = nn_module.path # pyrefly: ignore[missing-attribute] # We modify the path of the current nn module in place. This is a little # bit hacky but should be good as a temporary solution. - nn_module.scope.path += (name,) + nn_module.scope.path += (name,) # pyrefly: ignore[missing-attribute] try: return call_fn(*args, **kwargs) finally: - nn_module.scope.path = old_path + nn_module.scope.path = old_path # pyrefly: ignore[missing-attribute] return wrapped @@ -556,7 +556,7 @@ def __getattr__(self, name: str): maybe_method = getattr(self.nnx_class, name, None) if callable(maybe_method): method = partial(self.__call__, nnx_method=maybe_method) - method.__self__ = self + method.__self__ = self # pyrefly: ignore[missing-attribute] return method return super().__getattribute__(name) @@ -701,12 +701,12 @@ class ToLinenPartial(ToLinen): def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) - cls.__init__ = __init__ + cls.__init__ = __init__ # pyrefly: ignore[bad-assignment] ToLinenPartial.__name__ = class_name ToLinenPartial.__qualname__ = class_name - ToLinenPartial.__init__ = __init__ + ToLinenPartial.__init__ = __init__ # pyrefly: ignore[bad-assignment] ToLinenPartial.module_class = base_nnx_class return ToLinenPartial diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index 7174cf763a..0e538c9a6a 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -103,22 +103,22 @@ def _rhs_axis_metadata_wrapper( if len(x.shape) == 1: return nn.with_logical_partitioning((lambda: x), tuple(None for _ in mesh_axes))() - mesh_axes = list(mesh_axes) + mesh_axes = list(mesh_axes) # pyrefly: ignore[bad-assignment] if is_tiled: # tile_map is a mapping between original rank and a list of new, tiled rank. if len(mesh_axes) < len(tile_map): - mesh_axes = [None] * (len(tile_map) - len(mesh_axes)) + mesh_axes + mesh_axes = [None] * (len(tile_map) - len(mesh_axes)) + mesh_axes # pyrefly: ignore[unsupported-operation] new_mesh_axes = [None] * len(x.shape) for orig_rank, new_rank in tile_map.items(): assert new_rank assert len(new_rank) <= 2 new_mesh_axes[new_rank[-1]] = mesh_axes[orig_rank] - mesh_axes = new_mesh_axes + mesh_axes = new_mesh_axes # pyrefly: ignore[bad-assignment] if mesh_axes is not None and len(mesh_axes) > 0: for no_shard_idx in no_sharding_axis: if no_shard_idx < len(mesh_axes): - mesh_axes[no_shard_idx] = None + mesh_axes[no_shard_idx] = None # pyrefly: ignore[unsupported-operation] return nn.with_logical_partitioning((lambda: x), mesh_axes)() @@ -137,13 +137,13 @@ def _get_mixed_precision_cfg(self): is_tiled = False tiling_fn = None # pylint: disable=protected-access - module_path = "/".join(nn.module._context.module_stack[-1].path) + module_path = "/".join(nn.module._context.module_stack[-1].path) # pyrefly: ignore[missing-attribute] tile_size = -1 - for layer_name_re, layer_quant_dg in self.quant_dg.items(): + for layer_name_re, layer_quant_dg in self.quant_dg.items(): # pyrefly: ignore[missing-attribute] if re.fullmatch(layer_name_re, module_path): quant_dg, tile_size = layer_quant_dg if quant_dg is None: - quant_dg, tile_size = self.quant_dg[DEFAULT] + quant_dg, tile_size = self.quant_dg[DEFAULT] # pyrefly: ignore[bad-index] if tile_size != -1: is_tiled = True tiling_fn = functools.partial(_tiling_fn, tile_size=tile_size) @@ -272,7 +272,7 @@ def __call__( out_sharding=None, ) -> jax.Array: def custom_dot_general(*args, **kwargs): - return dot_general_qt.dot_general_qt(*args[:3], self.config) + return dot_general_qt.dot_general_qt(*args[:3], self.config) # pyrefly: ignore[bad-argument-count] with jax.disable_jit(): return jnp.einsum( @@ -392,7 +392,7 @@ def _get_int8_quant_config(config): if config.quantization_local_shard_count != 0: drhs_bits = 8 drhs_accumulator_dtype = jnp.int32 - drhs_local_aqt = aqt_config.LocalAqt(contraction_axis_shard_count=config.quantization_local_shard_count) + drhs_local_aqt = aqt_config.LocalAqt(contraction_axis_shard_count=config.quantization_local_shard_count) # pyrefly: ignore[unexpected-keyword] return aqt_config.config_v3( fwd_bits=8, dlhs_bits=8, @@ -430,30 +430,30 @@ def _build_const_scale_config( The AQT dot general config with constant scale config. """ if cst_bound_config.fwd_lhs_bound is not None: - aqt_dg.fwd.dg_quantizer.lhs.calibration = functools.partial( + aqt_dg.fwd.dg_quantizer.lhs.calibration = functools.partial( # pyrefly: ignore[missing-attribute] calibration.ConstantCalibration, bound=cst_bound_config.fwd_lhs_bound ) if cst_bound_config.fwd_rhs_bound is not None: - aqt_dg.fwd.dg_quantizer.rhs.calibration = functools.partial( + aqt_dg.fwd.dg_quantizer.rhs.calibration = functools.partial( # pyrefly: ignore[missing-attribute] calibration.ConstantCalibration, bound=cst_bound_config.fwd_rhs_bound ) if cst_bound_config.dlhs_lhs_bound: - aqt_dg.dlhs.dg_quantizer.lhs.calibration = functools.partial( + aqt_dg.dlhs.dg_quantizer.lhs.calibration = functools.partial( # pyrefly: ignore[missing-attribute] calibration.ConstantCalibration, bound=cst_bound_config.dlhs_lhs_bound ) if cst_bound_config.dlhs_rhs_bound is not None: - aqt_dg.dlhs.dg_quantizer.rhs.calibration = functools.partial( + aqt_dg.dlhs.dg_quantizer.rhs.calibration = functools.partial( # pyrefly: ignore[missing-attribute] calibration.ConstantCalibration, bound=cst_bound_config.dlhs_rhs_bound ) if cst_bound_config.drhs_lhs_bound is not None: - aqt_dg.drhs.dg_quantizer.lhs.calibration = functools.partial( + aqt_dg.drhs.dg_quantizer.lhs.calibration = functools.partial( # pyrefly: ignore[missing-attribute] calibration.ConstantCalibration, bound=cst_bound_config.drhs_lhs_bound ) if cst_bound_config.drhs_rhs_bound is not None: - aqt_dg.drhs.dg_quantizer.rhs.calibration = functools.partial( + aqt_dg.drhs.dg_quantizer.rhs.calibration = functools.partial( # pyrefly: ignore[missing-attribute] calibration.ConstantCalibration, bound=cst_bound_config.drhs_rhs_bound ) @@ -484,17 +484,17 @@ def _build_per_tensor_config( The AQT dot general config with per tensor config. """ if per_tensor_scales.fwd_lhs: - aqt_dg.fwd.dg_quantizer.lhs.calib_shared_axes = "per_tensor" + aqt_dg.fwd.dg_quantizer.lhs.calib_shared_axes = "per_tensor" # pyrefly: ignore[missing-attribute] if per_tensor_scales.fwd_rhs: - aqt_dg.fwd.dg_quantizer.rhs.calib_shared_axes = "per_tensor" + aqt_dg.fwd.dg_quantizer.rhs.calib_shared_axes = "per_tensor" # pyrefly: ignore[missing-attribute] if per_tensor_scales.dlhs_lhs: - aqt_dg.dlhs.dg_quantizer.lhs.calib_shared_axes = "per_tensor" + aqt_dg.dlhs.dg_quantizer.lhs.calib_shared_axes = "per_tensor" # pyrefly: ignore[missing-attribute] if per_tensor_scales.dlhs_rhs: - aqt_dg.dlhs.dg_quantizer.rhs.calib_shared_axes = "per_tensor" + aqt_dg.dlhs.dg_quantizer.rhs.calib_shared_axes = "per_tensor" # pyrefly: ignore[missing-attribute] if per_tensor_scales.drhs_lhs: - aqt_dg.drhs.dg_quantizer.lhs.calib_shared_axes = "per_tensor" + aqt_dg.drhs.dg_quantizer.lhs.calib_shared_axes = "per_tensor" # pyrefly: ignore[missing-attribute] if per_tensor_scales.drhs_rhs: - aqt_dg.drhs.dg_quantizer.rhs.calib_shared_axes = "per_tensor" + aqt_dg.drhs.dg_quantizer.rhs.calib_shared_axes = "per_tensor" # pyrefly: ignore[missing-attribute] return aqt_dg @@ -565,9 +565,9 @@ def _dot_general_make(quant_cfg): rhs_scale = quant_cfg[_W_SCALE] aqt_dg = aqt_config.dot_general_make(lhs_bits=lhs_bits, rhs_bits=rhs_bits) if lhs_scale < 1.0: - aqt_dg.fwd.dg_quantizer.lhs.calibration = functools.partial(calibration.AbsMaxCalibration, scale=lhs_scale) + aqt_dg.fwd.dg_quantizer.lhs.calibration = functools.partial(calibration.AbsMaxCalibration, scale=lhs_scale) # pyrefly: ignore[missing-attribute] if rhs_scale < 1.0: - aqt_dg.fwd.dg_quantizer.rhs.calibration = functools.partial(calibration.AbsMaxCalibration, scale=rhs_scale) + aqt_dg.fwd.dg_quantizer.rhs.calibration = functools.partial(calibration.AbsMaxCalibration, scale=rhs_scale) # pyrefly: ignore[missing-attribute] return aqt_dg @@ -728,7 +728,7 @@ def _apply_linen_module_in_nnx(linen_module_cls, op_id, *args, **kwargs): if is_nnx: attr_name = f"_qwix_fp8_gpu_{op_id}" - if not hasattr(parent, attr_name): + if not hasattr(parent, attr_name): # pyrefly: ignore[unbound-name] rngs = getattr(parent, "qwix_rngs", None) if rngs is None: parent_rngs = getattr(parent, "rngs", None) @@ -1038,7 +1038,7 @@ def generate_quantizer_set( variable_collection=OVERWRITE_WITH_GRADIENT, quantization_checkpoint_name="quantization", fp8_recipe=fp8_recipe, - n_groups=n_groups, + n_groups=n_groups, # pyrefly: ignore[bad-argument-type] ) @nn.compact diff --git a/src/maxtext/multimodal/processor.py b/src/maxtext/multimodal/processor.py index 381695915a..f15617bf10 100644 --- a/src/maxtext/multimodal/processor.py +++ b/src/maxtext/multimodal/processor.py @@ -149,15 +149,15 @@ def prepare_text_for_image_fusion(tokens, config, processor_output=None): if config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]: from maxtext.multimodal.processor_gemma3 import add_extra_tokens_for_images_gemma3 # pylint: disable=import-outside-toplevel - return add_extra_tokens_for_images_gemma3(tokens, max_num_images=processor_output.num_images) + return add_extra_tokens_for_images_gemma3(tokens, max_num_images=processor_output.num_images) # pyrefly: ignore[missing-attribute] elif config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]: from maxtext.multimodal.processor_gemma4 import add_extra_tokens_for_images_gemma4 # pylint: disable=import-outside-toplevel - return add_extra_tokens_for_images_gemma4(tokens, max_num_images=processor_output.num_images) + return add_extra_tokens_for_images_gemma4(tokens, max_num_images=processor_output.num_images) # pyrefly: ignore[missing-attribute] elif config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]: from maxtext.multimodal.processor_llama4 import add_extra_tokens_for_images_llama4 # pylint: disable=import-outside-toplevel - return add_extra_tokens_for_images_llama4(tokens, processor_output) + return add_extra_tokens_for_images_llama4(tokens, processor_output) # pyrefly: ignore[bad-argument-type] elif config.model_name in ["qwen3-omni-30b-a3b", "qwen3-vl-2b", "qwen3-vl-4b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]: from maxtext.multimodal.processor_qwen3_omni import add_extra_tokens_for_qwen3_omni # pylint: disable=import-outside-toplevel diff --git a/src/maxtext/multimodal/processor_gemma3.py b/src/maxtext/multimodal/processor_gemma3.py index 4b444fe935..7e059edd46 100644 --- a/src/maxtext/multimodal/processor_gemma3.py +++ b/src/maxtext/multimodal/processor_gemma3.py @@ -86,7 +86,7 @@ def preprocess_mm_data_gemma3(images): def get_image_offsets_gemma3(processor_output: mm_utils.PreprocessorOutput | None): """Get the increase in total token count after inserting image token placeholders""" has_images = processor_output is not None and processor_output.pixel_values is not None - num_images = processor_output.pixel_values.shape[0] if has_images else 1 + num_images = processor_output.pixel_values.shape[0] if has_images else 1 # pyrefly: ignore[missing-attribute] return ( GEMMA_NUM_TOKENS_PER_MEDIA - 1 ) * num_images # -1 because is already present in the input tokens. diff --git a/src/maxtext/multimodal/processor_gemma4.py b/src/maxtext/multimodal/processor_gemma4.py index 5671c750ac..e261eb8223 100644 --- a/src/maxtext/multimodal/processor_gemma4.py +++ b/src/maxtext/multimodal/processor_gemma4.py @@ -78,7 +78,7 @@ def preprocess_mm_data_gemma4(images): def get_image_offsets_gemma4(processor_output: mm_utils.PreprocessorOutput | None): """Gets the increase in total token count after inserting image token placeholders.""" has_images = processor_output is not None and processor_output.pixel_values is not None - num_images = processor_output.pixel_values.shape[0] if has_images else 1 + num_images = processor_output.pixel_values.shape[0] if has_images else 1 # pyrefly: ignore[missing-attribute] # Calculate soft tokens taking 3x3 pooling into account num_patches = (GEMMA4_IMAGE_HEIGHT // GEMMA4_PATCH_SIZE) * (GEMMA4_IMAGE_WIDTH // GEMMA4_PATCH_SIZE) diff --git a/src/maxtext/multimodal/processor_qwen3_omni.py b/src/maxtext/multimodal/processor_qwen3_omni.py index 787dd82011..dbf838cf2c 100644 --- a/src/maxtext/multimodal/processor_qwen3_omni.py +++ b/src/maxtext/multimodal/processor_qwen3_omni.py @@ -167,7 +167,7 @@ def maybe_pad_video_values_to_max_grid( if video_values.ndim != 5: raise ValueError(f"video_values must have shape (batch, channels, time, height, width), got {video_values.shape}.") - max_t, max_h, max_w = (int(dim) for dim in max_grid) + max_t, max_h, max_w = (int(dim) for dim in max_grid) # pyrefly: ignore[bad-argument-type] actual_t, actual_h, actual_w = (int(dim) for dim in video_grid_thw[0]) if actual_t > max_t or actual_h > max_h or actual_w > max_w: raise ValueError( @@ -348,6 +348,7 @@ def calculate_video_frame_range( # Validate frame order if start_frame >= end_frame: raise ValueError( + # pyrefly: ignore[unbound-name] f"Invalid time range: Start frame {start_frame} (at {video_start_clamped if video_start is not None else 0}s) " f"exceeds end frame {end_frame} (at {video_end_clamped if video_end is not None else max_duration}s). " f"Video duration: {max_duration:.2f}s ({total_frames} frames @ {video_fps}fps)" @@ -1112,7 +1113,7 @@ def get_rope_index( # Process modality-specific content # Audio Only if min_ed == ed_audio_start: - audio_len = _get_feat_extract_output_lengths(audio_lengths[audio_idx]).item() + audio_len = _get_feat_extract_output_lengths(audio_lengths[audio_idx]).item() # pyrefly: ignore[unsupported-operation] audio_pos = np.arange(audio_len).reshape(1, -1).repeat(3, axis=0) + st_idx llm_pos_ids_list.append(audio_pos) @@ -1122,45 +1123,45 @@ def get_rope_index( # Image Only elif min_ed == ed_vision_start and input_tokens[ed_vision_start + 1] == qwen_tokens.image_pad: - grid_t = image_grid_thw[image_idx, 0].item() - grid_hs = image_grid_thw[:, 1] - grid_ws = image_grid_thw[:, 2] + grid_t = image_grid_thw[image_idx, 0].item() # pyrefly: ignore[unsupported-operation] + grid_hs = image_grid_thw[:, 1] # pyrefly: ignore[unsupported-operation] + grid_ws = image_grid_thw[:, 2] # pyrefly: ignore[unsupported-operation] t_index = np.arange(grid_t, dtype=np.float32) * 1 * position_id_per_seconds - image_pos = get_llm_pos_ids_for_vision(st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws) + image_pos = get_llm_pos_ids_for_vision(st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws) # pyrefly: ignore[bad-argument-type] llm_pos_ids_list.append(image_pos) - image_len = int(np.prod(image_grid_thw[image_idx]).item() // (spatial_merge_size**2)) + image_len = int(np.prod(image_grid_thw[image_idx]).item() // (spatial_merge_size**2)) # pyrefly: ignore[unsupported-operation] st += int(text_len + bos_len + image_len + eos_len) image_idx += 1 remain_images -= 1 # Video Only elif min_ed == ed_vision_start and input_tokens[ed_vision_start + 1] == qwen_tokens.video_pad: - grid_t = video_grid_thw[video_idx, 0].item() - grid_hs = video_grid_thw[:, 1] - grid_ws = video_grid_thw[:, 2] - t_index = np.arange(grid_t, dtype=np.float32) * second_per_grids[video_idx].item() * position_id_per_seconds + grid_t = video_grid_thw[video_idx, 0].item() # pyrefly: ignore[unsupported-operation] + grid_hs = video_grid_thw[:, 1] # pyrefly: ignore[unsupported-operation] + grid_ws = video_grid_thw[:, 2] # pyrefly: ignore[unsupported-operation] + t_index = np.arange(grid_t, dtype=np.float32) * second_per_grids[video_idx].item() * position_id_per_seconds # pyrefly: ignore[unsupported-operation] - video_pos = get_llm_pos_ids_for_vision(st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws) + video_pos = get_llm_pos_ids_for_vision(st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws) # pyrefly: ignore[bad-argument-type] llm_pos_ids_list.append(video_pos) - video_len = int(np.prod(video_grid_thw[video_idx]).item() // (spatial_merge_size**2)) + video_len = int(np.prod(video_grid_thw[video_idx]).item() // (spatial_merge_size**2)) # pyrefly: ignore[unsupported-operation] st += int(text_len + bos_len + video_len + eos_len) video_idx += 1 remain_videos -= 1 # Audio in Video (interleaved) elif min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start: - audio_len = _get_feat_extract_output_lengths(audio_lengths[audio_idx]).item() + audio_len = _get_feat_extract_output_lengths(audio_lengths[audio_idx]).item() # pyrefly: ignore[unsupported-operation] audio_llm_pos_ids = np.arange(audio_len).reshape(1, -1).repeat(3, axis=0) + st_idx - grid_t = video_grid_thw[video_idx, 0].item() - grid_hs = video_grid_thw[:, 1] - grid_ws = video_grid_thw[:, 2] - t_index = np.arange(grid_t, dtype=np.float32) * second_per_grids[video_idx].item() * position_id_per_seconds + grid_t = video_grid_thw[video_idx, 0].item() # pyrefly: ignore[unsupported-operation] + grid_hs = video_grid_thw[:, 1] # pyrefly: ignore[unsupported-operation] + grid_ws = video_grid_thw[:, 2] # pyrefly: ignore[unsupported-operation] + t_index = np.arange(grid_t, dtype=np.float32) * second_per_grids[video_idx].item() * position_id_per_seconds # pyrefly: ignore[unsupported-operation] - video_llm_pos_ids = get_llm_pos_ids_for_vision(st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws) + video_llm_pos_ids = get_llm_pos_ids_for_vision(st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws) # pyrefly: ignore[bad-argument-type] # Interleave audio and video based on temporal ordering video_data_index = 0 @@ -1179,7 +1180,7 @@ def get_rope_index( if audio_data_index < audio_llm_pos_ids.shape[1]: llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index:]) - video_len = int(np.prod(video_grid_thw[video_idx]).item() // (spatial_merge_size**2)) + video_len = int(np.prod(video_grid_thw[video_idx]).item() // (spatial_merge_size**2)) # pyrefly: ignore[unsupported-operation] st += int(text_len + bos_len + audio_len + video_len + eos_len) audio_idx += 1 diff --git a/src/maxtext/multimodal/utils.py b/src/maxtext/multimodal/utils.py index 65b5670fc1..a068cc6fa9 100644 --- a/src/maxtext/multimodal/utils.py +++ b/src/maxtext/multimodal/utils.py @@ -188,7 +188,7 @@ def _merge_mm_embeddings_inner( # Find positions in the text sequence to place the multimodal embeddings. # The `size` argument ensures a fixed shape for JIT compilation. - target_pos = jnp.nonzero(mask, size=multimodal_embeddings.shape[0]) + target_pos = jnp.nonzero(mask, size=multimodal_embeddings.shape[0]) # pyrefly: ignore[bad-argument-type] target_pos = target_pos[0] # jnp.nonzero returns a tuple of arrays # Save the embedding at the first position. @@ -452,7 +452,7 @@ def spectrogram( # center pad the waveform if center: padding = [(int(frame_length // 2), int(frame_length // 2))] - waveform = np.pad(waveform, padding, mode=pad_mode) + waveform = np.pad(waveform, padding, mode=pad_mode) # pyrefly: ignore[no-matching-overload] # promote to float64, since np.fft uses float64 internally waveform = waveform.astype(np.float64) @@ -545,7 +545,7 @@ def hertz_to_mel(freq: Union[float, np.ndarray], mel_scale: str = "htk") -> Unio if isinstance(freq, np.ndarray): log_region = freq >= min_log_hertz - mels[log_region] = min_log_mel + np.log(freq[log_region] / min_log_hertz) * logstep + mels[log_region] = min_log_mel + np.log(freq[log_region] / min_log_hertz) * logstep # pyrefly: ignore[unsupported-operation] elif freq >= min_log_hertz: mels = min_log_mel + np.log(freq / min_log_hertz) * logstep @@ -603,7 +603,7 @@ def mel_to_hertz(mels: Union[float, np.ndarray], mel_scale: str = "htk") -> Unio if isinstance(mels, np.ndarray): log_region = mels >= min_log_mel - freq[log_region] = min_log_hertz * np.exp(logstep * (mels[log_region] - min_log_mel)) + freq[log_region] = min_log_hertz * np.exp(logstep * (mels[log_region] - min_log_mel)) # pyrefly: ignore[unsupported-operation] elif mels >= min_log_mel: freq = min_log_hertz * np.exp(logstep * (mels - min_log_mel)) @@ -688,11 +688,11 @@ def mel_filter_bank( # frequencies of FFT bins in Hz fft_freqs = np.linspace(0, sampling_rate // 2, num_frequency_bins) - mel_filters = _create_triangular_filter_bank(fft_freqs, filter_freqs) + mel_filters = _create_triangular_filter_bank(fft_freqs, filter_freqs) # pyrefly: ignore[bad-argument-type] if norm is not None and norm == "slaney": # Slaney-style mel is scaled to be approx constant energy per channel - enorm = 2.0 / (filter_freqs[2 : num_mel_filters + 2] - filter_freqs[:num_mel_filters]) + enorm = 2.0 / (filter_freqs[2 : num_mel_filters + 2] - filter_freqs[:num_mel_filters]) # pyrefly: ignore[bad-index] mel_filters *= np.expand_dims(enorm, 0) if (mel_filters.max(axis=0) == 0.0).any(): diff --git a/src/maxtext/trainers/diloco/diloco.py b/src/maxtext/trainers/diloco/diloco.py index 1b7253c41d..48accf4779 100644 --- a/src/maxtext/trainers/diloco/diloco.py +++ b/src/maxtext/trainers/diloco/diloco.py @@ -157,11 +157,11 @@ def add_diloco_dim(x): # for Linen under abstract_state.params. if config.pure_nnx: _, model_params, _ = nnx.split(abstract_state.model, nnx.Param, ...) - model_params = model_params.to_pure_dict() + model_params = model_params.to_pure_dict() # pyrefly: ignore[missing-attribute] _, model_params_sharding, _ = nnx.split( state_mesh_shardings.model, nnx.Param, ... ) - model_params_sharding = model_params_sharding.to_pure_dict() + model_params_sharding = model_params_sharding.to_pure_dict() # pyrefly: ignore[missing-attribute] else: model_params = abstract_state.params model_params_sharding = state_mesh_shardings.params @@ -175,7 +175,7 @@ def add_diloco_dim(x): inner_state=inner_state, params=model_params, outer_opt_state=outer_opt_state, - step=abstract_step, + step=abstract_step, # pyrefly: ignore[bad-argument-type] ) # Build shardings @@ -190,7 +190,7 @@ def add_diloco_dim(x): inner_state=inner_state_shardings, params=model_params_sharding, outer_opt_state=outer_opt_state_sharding, - step=None, + step=None, # pyrefly: ignore[bad-argument-type] ) return diloco_state, diloco_state_shardings, inner_state_shardings @@ -221,7 +221,7 @@ def init_diloco_state() -> tuple[DiLoCoTrainState, PyTree]: # for Linen under state.params. if config.pure_nnx: _, outer_params, _ = nnx.split(state.model, nnx.Param, ...) - outer_params = outer_params.to_pure_dict() + outer_params = outer_params.to_pure_dict() # pyrefly: ignore[missing-attribute] else: outer_params = state.params outer_opt_state = outer_optimizer.init(outer_params) @@ -269,7 +269,7 @@ def synchronize(state): _, inner_model_params, _ = nnx.split( state.inner_state.model, nnx.Param, ... ) - inner_model_params = inner_model_params.to_pure_dict() + inner_model_params = inner_model_params.to_pure_dict() # pyrefly: ignore[missing-attribute] else: inner_model_params = state.inner_state.params model_delta = jax.tree.map(lambda x, y: y - x, inner_model_params, broadcast_outer_params) diff --git a/src/maxtext/trainers/post_train/distillation/distillation_utils.py b/src/maxtext/trainers/post_train/distillation/distillation_utils.py index f063cdb23a..1946c2f2e8 100644 --- a/src/maxtext/trainers/post_train/distillation/distillation_utils.py +++ b/src/maxtext/trainers/post_train/distillation/distillation_utils.py @@ -508,7 +508,7 @@ def compute_loss( log_t_p_T_sparse = jax.nn.log_softmax(t_logits / temperature, axis=-1) # 2. Gather Student unnormalized logits at the Teacher's exact Top-K indices - s_logits_sparse = jnp.take_along_axis(s_logits, teacher_output.top_k_indices, axis=-1) + s_logits_sparse = jnp.take_along_axis(s_logits, teacher_output.top_k_indices, axis=-1) # pyrefly: ignore[bad-argument-type] # 3. Normalize Student probabilities only over the exact same Top-K subset log_s_T_sparse = jax.nn.log_softmax(s_logits_sparse / temperature, axis=-1) @@ -558,7 +558,7 @@ def compute_loss( s_features_sliced = s_features_sliced.astype(jnp.float32) t_features_sliced = t_features_sliced.astype(jnp.float32) - feature_loss = beta_feature * self.feature_loss_fn(s_features_sliced, t_features_sliced, mask) + feature_loss = beta_feature * self.feature_loss_fn(s_features_sliced, t_features_sliced, mask) # pyrefly: ignore[not-callable] total_loss = base_logit_loss + feature_loss @@ -744,7 +744,7 @@ def save( grain_iters_to_save.append((local_iter, process_index, process_count_total)) # Use GrainCheckpointSave wrapper - cp_save_args["iter"] = GrainCheckpointSave(item=grain_iters_to_save) + cp_save_args["iter"] = GrainCheckpointSave(item=grain_iters_to_save) # pyrefly: ignore[bad-assignment] return self._checkpoint_manager.save( step, @@ -753,7 +753,7 @@ def save( force=force, ) - def maybe_restore( + def maybe_restore( # pyrefly: ignore[bad-override] self, model: Any, optimizer: Any = None, @@ -772,7 +772,7 @@ def maybe_restore( target_model = getattr(model, "student_model", model) step, _ = super().maybe_restore( - model=target_model, + model=target_model, # pyrefly: ignore[bad-argument-type] optimizer=optimizer, restore_only_lora_params=restore_only_lora_params, ) diff --git a/src/maxtext/trainers/post_train/distillation/train_distill.py b/src/maxtext/trainers/post_train/distillation/train_distill.py index 15197d1af4..952811bd37 100644 --- a/src/maxtext/trainers/post_train/distillation/train_distill.py +++ b/src/maxtext/trainers/post_train/distillation/train_distill.py @@ -203,7 +203,7 @@ def call_student(self, *args, **kwargs): return self.student_model(*args, **kwargs) def call_teacher(self, *args, **kwargs): - return jax.lax.stop_gradient(self.teacher_model(*args, **kwargs)) + return jax.lax.stop_gradient(self.teacher_model(*args, **kwargs)) # pyrefly: ignore[not-callable] class MaxTextDistillationTrainer(peft_trainer.PeftTrainer): @@ -232,7 +232,7 @@ def __init__( super().__init__(model=model, optimizer=dummy_optimizer, training_config=training_config, **kwargs) self.strategy = strategy - self.checkpoint_manager: distillation_utils.MaxTextCheckpointManager = None + self.checkpoint_manager: distillation_utils.MaxTextCheckpointManager = None # pyrefly: ignore[bad-assignment] # Per-step per-device TFLOPs (constants for the run): student fwd+bwd + teacher fwd-only. ( @@ -385,11 +385,12 @@ def _log_metrics(self, loss, step=None, additional_metrics=None, **kwargs): } ) for name, value in tflops_metrics.items(): - self.metrics_logger.log(self.metrics_prefix, name, value, self._mode, step) + self.metrics_logger.log(self.metrics_prefix, name, value, self._mode, step) # pyrefly: ignore[missing-attribute] # Console summary — keep it tight; everything else is in TensorBoard. if tflops_per_sec is not None and self._mode == metrics_logger.Mode.TRAIN: max_logging.log( + # pyrefly: ignore[unsupported-operation] f"step {step} | step_time={step_time_delta:.2f}s | " f"TFLOPs/s/device: {tflops_per_sec:.2f} " f"(student={self._tflops_student / step_time_delta:.2f}, " @@ -438,9 +439,9 @@ def _post_process_train_step(self, aux: dict[str, tuple[jax.Array, jax.Array]]) for name, value in aux.items(): if name not in self._buffered_train_metrics.additional_metrics: - self._buffered_train_metrics.additional_metrics[name] = ([], distillation_utils.weighted_mean) + self._buffered_train_metrics.additional_metrics[name] = ([], distillation_utils.weighted_mean) # pyrefly: ignore[unsupported-operation] - self._buffered_train_metrics.additional_metrics[name][0].append(value) + self._buffered_train_metrics.additional_metrics[name][0].append(value) # pyrefly: ignore[bad-argument-type] # Compact per-step summary: only the metrics that change run-to-run, and # only those that are nonzero (so MoE / feature-distill terms hide when off). diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index e7f5ea84b6..a9f3ea82b4 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -147,7 +147,7 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr encoder_images=data["images"] if config.use_multimodal else None, encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None, enable_dropout=config.enable_dropout if is_train else False, - rngs={"dropout": rng1, "params": aqt_rng}, + rngs={"dropout": rng1, "params": aqt_rng}, # pyrefly: ignore[bad-argument-type] mutable=mutable_collections, decoder_target_tokens=data["targets"], decoder_target_mask=data["targets_segmentation"], @@ -375,11 +375,11 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat if config.shard_optimizer_over_data: params = jax.tree.map( functools.partial(sharding.maybe_shard_with_name, shard_mode=config.shard_mode), - params, + params, # pyrefly: ignore[unbound-name] params_shardings, ) sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m - pure_params = params["params"] if sparsity_enabled else params + pure_params = params["params"] if sparsity_enabled else params # pyrefly: ignore[unbound-name] batch_stats = params.get("batch_stats", {}) grad_func = jax.value_and_grad(loss_fn, argnums=4, has_aux=True) diff --git a/src/maxtext/trainers/pre_train/train_compile.py b/src/maxtext/trainers/pre_train/train_compile.py index d472bb68d4..aa33ce73cb 100644 --- a/src/maxtext/trainers/pre_train/train_compile.py +++ b/src/maxtext/trainers/pre_train/train_compile.py @@ -187,7 +187,7 @@ def create_train_state_fn(): # Collect NNX activation shardings via an abstract forward pass (must run # after get_abstract_state, which only traces __init__). if config.debug_sharding and config.pure_nnx: - _collect_nnx_activation_shardings(_create_model_partial, config, topology_mesh) + _collect_nnx_activation_shardings(_create_model_partial, config, topology_mesh) # pyrefly: ignore[unbound-name] return ( shaped_train_args, diff --git a/src/maxtext/trainers/tokenizer/train_tokenizer.py b/src/maxtext/trainers/tokenizer/train_tokenizer.py index d3b4f35fce..58c63955fb 100644 --- a/src/maxtext/trainers/tokenizer/train_tokenizer.py +++ b/src/maxtext/trainers/tokenizer/train_tokenizer.py @@ -82,7 +82,7 @@ def build_grain_iterator(data_file_pattern: str, data_file_type: str, data_keys: if data_file_type == "parquet": dataset = grain.MapDataset.source(data_files) dataset = dataset.map(grain.experimental.ParquetIterDataset) - dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=len(data_files)) + dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=len(data_files)) # pyrefly: ignore[bad-argument-type] dataset = dataset.map(input_pipeline_utils.KeepFeatures(feature_names=list(data_keys))) return iter(dataset) elif data_file_type == "arrayrecord": @@ -94,7 +94,7 @@ def build_grain_iterator(data_file_pattern: str, data_file_type: str, data_keys: elif data_file_type == "tfrecord": dataset = grain.MapDataset.source(data_files) dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset) - dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=len(data_files)) + dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=len(data_files)) # pyrefly: ignore[bad-argument-type] dataset = dataset.map(input_pipeline_utils.ParseFeatures(list(data_keys), tokenize=True)) dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(list(data_keys), tokenize=True)) return iter(dataset) diff --git a/tests/integration/smoke/train_gpu_smoke_test.py b/tests/integration/smoke/train_gpu_smoke_test.py index 1df3ef630b..40dbf8ec35 100644 --- a/tests/integration/smoke/train_gpu_smoke_test.py +++ b/tests/integration/smoke/train_gpu_smoke_test.py @@ -41,7 +41,7 @@ def setUp(self): def test_tiny_config(self): test_tmpdir = os.environ.get("TEST_TMPDIR") # pylint: disable=unused-variable train_main( - [ + [ # pyrefly: ignore[bad-argument-type] None, get_test_config_path("gpu/gpu_smoke_test.yml"), # pylint: disable=f-string-without-interpolation diff --git a/tests/integration/smoke/train_int8_smoke_test.py b/tests/integration/smoke/train_int8_smoke_test.py index 274d3a124c..f31f6b976c 100644 --- a/tests/integration/smoke/train_int8_smoke_test.py +++ b/tests/integration/smoke/train_int8_smoke_test.py @@ -40,7 +40,7 @@ def setUp(self): def test_tiny_config(self): test_tmpdir = os.environ.get("TEST_TMPDIR") # pylint: disable=unused-variable train_main( - [ + [ # pyrefly: ignore[bad-argument-type] None, get_test_config_path(), # pylint: disable=f-string-without-interpolation diff --git a/tests/integration/smoke/train_using_ragged_dot_smoke_test.py b/tests/integration/smoke/train_using_ragged_dot_smoke_test.py index 8a3b7b5bcb..6af73345db 100644 --- a/tests/integration/smoke/train_using_ragged_dot_smoke_test.py +++ b/tests/integration/smoke/train_using_ragged_dot_smoke_test.py @@ -37,7 +37,7 @@ class Train(parameterized.TestCase): def test_tiny_config(self, quantization: str): test_tmpdir = os.environ.get("TEST_TMPDIR", gettempdir()) outputs_dir = os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR", test_tmpdir) - train_main([ + train_main([ # pyrefly: ignore[bad-argument-type] None, get_test_config_path(), f"base_output_directory={test_tmpdir}", diff --git a/tests/integration/sparsity_test.py b/tests/integration/sparsity_test.py index 10dd93d7c4..7154ed00f8 100644 --- a/tests/integration/sparsity_test.py +++ b/tests/integration/sparsity_test.py @@ -91,7 +91,7 @@ def test_different_quant_sparsity_configs(self, quantization: str, use_sparsity: "weight_sparsity_update_step=1", ] ) - train_main(args) + train_main(args) # pyrefly: ignore[bad-argument-type] if __name__ == "__main__": diff --git a/tests/unit/elastic_utils_test.py b/tests/unit/elastic_utils_test.py index 5c344feb80..e813baaa7c 100644 --- a/tests/unit/elastic_utils_test.py +++ b/tests/unit/elastic_utils_test.py @@ -79,7 +79,7 @@ def setUp(self): elastic_utils.max_logging = self.fake_logging # Hook up pathwaysutils.elastic.manager.Manager to return our fake_manager - pathwaysutils.elastic.manager.Manager = lambda *args, **kwargs: self.fake_manager + pathwaysutils.elastic.manager.Manager = lambda *args, **kwargs: self.fake_manager # pyrefly: ignore[bad-assignment] pathwaysutils.elastic.manager.ScaleUpSignalError = ScaleUpSignalError # Reset global state for testing is no longer needed @@ -91,7 +91,7 @@ def tearDown(self): elastic_utils.gcs_utils = self.original_gcs_utils elastic_utils.max_logging = self.original_max_logging pathwaysutils.elastic.manager.Manager = self.original_manager_class - pathwaysutils.elastic.manager.ScaleUpSignalError = self.original_scale_up_signal_error + pathwaysutils.elastic.manager.ScaleUpSignalError = self.original_scale_up_signal_error # pyrefly: ignore[bad-assignment] elastic_utils.elastic_manager = None elastic_utils.pending_reinit_recorder = None elastic_utils.pending_elastic_event_type = None @@ -355,7 +355,7 @@ def test_record_elastic_wait_end_and_reinit_start_noop_on_first_attempt(self): def test_record_elastic_wait_end_and_reinit_start(self): """Test recording end of slice down and start of reinit.""" - elastic_utils.pending_elastic_event_type = "elastic_slice_down" + elastic_utils.pending_elastic_event_type = "elastic_slice_down" # pyrefly: ignore[bad-assignment] fake_recorder = Mock() elastic_utils.record_elastic_wait_end_and_reinit_start(fake_recorder) diff --git a/tests/utils/attention_test_util.py b/tests/utils/attention_test_util.py index 29777d9549..12ff18591a 100644 --- a/tests/utils/attention_test_util.py +++ b/tests/utils/attention_test_util.py @@ -215,8 +215,8 @@ def forward_with_context_expert_parallelism( nn_partitioning.get_axis_rules(), ) pos_spec = nn_partitioning.logical_to_mesh_axes((None, length_axis), nn_partitioning.get_axis_rules()) - lnx_sharding = NamedSharding(mesh_cp, lnx_spec) - pos_sharding = NamedSharding(mesh_cp, pos_spec) + lnx_sharding = NamedSharding(mesh_cp, lnx_spec) # pyrefly: ignore[bad-argument-type] + pos_sharding = NamedSharding(mesh_cp, pos_spec) # pyrefly: ignore[bad-argument-type] lnx = jax.device_put(lnx, lnx_sharding) decoder_segment_ids = jax.device_put(decoder_segment_ids, pos_sharding)