diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 69d8dc0f6..c5f2f118f 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -123,6 +123,10 @@ grad_dtype: "float32" # activation dtypes. dtype: "bfloat16" + +# GDN precision configuration +gdn_state_dtype: "float32" +gdn_decay_dtype: "float32" # used to configure quantization in the transformer layers, defaults to null implying bf16. # possible alternative settings are as follows: # 'int8' for dynamic range quantization using 8-bits @@ -1357,6 +1361,8 @@ gdn_chunk_size: 64 use_qk_norm_in_gdn: true # The ratio of dimension to apply ROPE on partial_rotary_factor: 1.0 +# Whether to use GDN Pallas kernel +use_gdn_kernel: false use_tokamax_splash: false # Setting this flag will use a non-pallas implementation. diff --git a/src/maxtext/configs/pyconfig.py b/src/maxtext/configs/pyconfig.py index f784f27a6..b5bd17ec0 100644 --- a/src/maxtext/configs/pyconfig.py +++ b/src/maxtext/configs/pyconfig.py @@ -337,6 +337,8 @@ def __init__(self, pydantic_config: types.MaxTextConfig): final_dict["dtype"] = jnp.dtype(final_dict["dtype"]) final_dict["grad_dtype"] = jnp.dtype(final_dict["grad_dtype"]) final_dict["weight_dtype"] = jnp.dtype(final_dict["weight_dtype"]) + final_dict["gdn_state_dtype"] = jnp.dtype(final_dict.get("gdn_state_dtype", "float32")) + final_dict["gdn_decay_dtype"] = jnp.dtype(final_dict.get("gdn_decay_dtype", "float32")) final_dict["mu_dtype"] = ( final_dict["weight_dtype"] if not final_dict["mu_dtype"] else jnp.dtype(final_dict["mu_dtype"]) ) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index d35d6b9be..0d1686aa2 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -480,6 +480,11 @@ class DataTypes(BaseModel): description="If True, sets activations to float32 before the nonlinearity.", ) dtype_mm: str = Field("float32", description="Data type for multimodal model's vision encoder") + gdn_state_dtype: DType = Field(DType.FLOAT32, description="The data type for GDN recurrent states.") + gdn_decay_dtype: DType = Field( + DType.FLOAT32, + description="The data type for GDN decay parameters (A_log, dt_bias).", + ) class Quantization(BaseModel): @@ -1154,6 +1159,10 @@ class Qwen3Next(BaseModel): description="Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.", ) partial_rotary_factor: float = Field(1.0, description="The ratio of dimension to apply ROPE on") + use_gdn_kernel: bool = Field( + False, + description="Whether to use GDN Pallas kernel.", + ) # ---------------------------------------------------------------------------- diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 24da5a703..e28343465 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -30,7 +30,7 @@ from jax.experimental import xla_metadata import jax.nn import jax.numpy as jnp -from jax.sharding import Mesh +from jax.sharding import Mesh, PartitionSpec as P from maxtext.common.common_types import Array, AttentionType, BATCH, Config, DType, EMBED, LENGTH, MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_TRAIN from maxtext.common.common_types import KV_BATCH, KV_HEAD, ShardMode from maxtext.inference import kvcache @@ -581,8 +581,8 @@ def a_log_init(key, shape, dtype=jnp.float32): a_vals = jax.random.uniform(key, shape=shape, dtype=dtype, minval=1e-9, maxval=16.0) return jnp.log(a_vals) - self.A_log = nnx.Param(a_log_init(rngs.params(), (self.num_v_heads,), dtype=cfg.weight_dtype)) - self.dt_bias = nnx.Param(nnx.initializers.ones(rngs.params(), (self.num_v_heads,), dtype=cfg.weight_dtype)) + self.A_log = nnx.Param(a_log_init(rngs.params(), (self.num_v_heads,), dtype=jnp.float32)) + self.dt_bias = nnx.Param(nnx.initializers.ones(rngs.params(), (self.num_v_heads,), dtype=jnp.float32)) self.norm = Qwen3NextRMSNormGated( num_features=self.head_v_dim, # Normalize over the head dimension (D_v) @@ -653,6 +653,12 @@ def __call__( # hidden_states: (B, S, E) cfg = self.config batch, seq_len, _ = hidden_states.shape + decay_dtype = getattr(cfg, "gdn_decay_dtype", jnp.float32) + if isinstance(decay_dtype, str): + decay_dtype = getattr(jnp, decay_dtype, jnp.float32) + state_dtype = getattr(cfg, "gdn_state_dtype", jnp.float32) + if isinstance(state_dtype, str): + state_dtype = getattr(jnp, state_dtype, jnp.float32) flat_sharding, head_sharding, state_sharding = self._explicit_activation_shardings(batch) active_cache = kv_cache if kv_cache is not None else self.cache @@ -772,7 +778,6 @@ def __call__( truncate_sharded_tensor, ) from tpu_inference.utils import get_mesh_shape_product # pylint: disable=import-outside-toplevel # pytype: disable=import-error - from jax.sharding import PartitionSpec as P_spec # pylint: disable=import-outside-toplevel # pytype: disable=import-error except ImportError as e: raise ImportError( "GDN attention kernel require the vllm-tpu package. Please install it with `pip install vllm-tpu`." @@ -795,8 +800,8 @@ def __call__( mixed_qkv = jax.shard_map( lambda q, k, v: jnp.concatenate([q, k, v], axis=-1), mesh=self.mesh, - in_specs=(P_spec(attn_data, attn_head),) * 3, - out_specs=P_spec(attn_data, attn_head), + in_specs=(P(attn_data, attn_head),) * 3, + out_specs=P(attn_data, attn_head), check_vma=False, )(q_flat, k_flat, v_flat) @@ -840,8 +845,8 @@ def __call__( recurrent_state_paged, conv_weight, None, # conv_bias: MaxText conv1d uses use_bias=False. - jnp.asarray(self.A_log[...], dtype=cfg.dtype), - jnp.asarray(self.dt_bias[...], dtype=cfg.dtype), + jnp.asarray(self.A_log[...], dtype=decay_dtype), + jnp.asarray(self.dt_bias[...], dtype=decay_dtype), state_indices, query_start_loc, attention_metadata.request_distribution, # pyrefly: ignore[missing-attribute] @@ -918,178 +923,297 @@ def extract_state(c_in, v_len): else: conv_input = jnp.pad(qkv, ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) - # Perform the convolution. - conv_out = self.conv1d(conv_input, out_sharding=flat_sharding) - # Slice the output to match the original input sequence length. - conv_out = conv_out[:, -seq_len:, :] - qkv_conv = jax.nn.silu(conv_out.astype(jnp.float32)).astype(cfg.dtype) - # q_conv shape: (B, S, key_dim), k_conv shape: (B, S, key_dim), v_conv shape: (B, S, value_dim) - q_conv, k_conv, v_conv = jnp.split(qkv_conv, [self.key_dim, 2 * self.key_dim], axis=-1) - - # Reshape for multi-head processing - # query shape: (B, S, H_k, D_k) - query = jnp.reshape( - q_conv, - (batch, seq_len, self.num_k_heads, self.head_k_dim), - out_sharding=head_sharding, - ) - # key shape: (B, S, H_k, D_k) - key = jnp.reshape( - k_conv, - (batch, seq_len, self.num_k_heads, self.head_k_dim), - out_sharding=head_sharding, - ) - # value shape: (B, S, H_v, D_v) - value = jnp.reshape( - v_conv, - (batch, seq_len, self.num_v_heads, self.head_v_dim), - out_sharding=head_sharding, - ) - - # ========================================================================= - # STEP C: Gated Delta Rule Recurrence - # ========================================================================= - A_log = jnp.asarray(self.A_log[...], dtype=cfg.dtype) - dt_bias = jnp.asarray(self.dt_bias[...], dtype=cfg.dtype) - if cfg.shard_mode == ShardMode.EXPLICIT: - # Both are stored replicated but broadcast against (B, S, H_v) activations whose - # head axis is sharded, and explicit sharding requires broadcast operands to - # agree -- the same fix `_align_scale_with_normalized_axis` applies to the norm - # scales. - head_spec = jax.sharding.PartitionSpec(jax.typeof(a).sharding.spec[-1]) - A_log = jax.sharding.reshard(A_log, head_spec) - dt_bias = jax.sharding.reshard(dt_bias, head_spec) - # beta shape: (B, S, H_v) - beta = jax.nn.sigmoid(b) - # g shape: (B, S, H_v) - g = -jnp.exp(A_log) * jax.nn.softplus(a + dt_bias) - - if decoder_segment_ids is not None: - mask = decoder_segment_ids != 0 - # Apply mask by broadcasting to respective shapes - key = jnp.where(mask[..., None, None], key, 0.0) - value = jnp.where(mask[..., None, None], value, 0.0) - g = jnp.where(mask[..., None], g, 0.0) - - if self.num_v_heads > self.num_k_heads and self.num_v_heads % self.num_k_heads == 0: - repeats = self.num_v_heads // self.num_k_heads - # query shape after repeat: (B, S, H_v, D_k) - query = jnp.repeat(query, repeats, axis=2, out_sharding=head_sharding) - # key shape after repeat: (B, S, H_v, D_k) - key = jnp.repeat(key, repeats, axis=2, out_sharding=head_sharding) - - if seq_len == 1 and model_mode == MODEL_MODE_AUTOREGRESSIVE: - core_attn_out, next_recurrent_state = jax_ar_gated_delta_rule( - query, - key, - value, - g, - beta, - initial_state=recurrent_state, # pyrefly: ignore[bad-argument-type] - use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, - compute_dtype=cfg.dtype, + if getattr(cfg, "use_gdn_kernel", False): + try: + from maxtext.models.kernels.gdn.gdn_bwd_pallas import gdn_decoupled_conv1d # pylint: disable=import-outside-toplevel + except ImportError: + try: + from .kernels.gdn.gdn_bwd_pallas import gdn_decoupled_conv1d # pylint: disable=import-outside-toplevel + except ImportError: + from .kernels.gdn import gdn_decoupled_conv1d # pylint: disable=import-outside-toplevel + + conv_state_arg = ( + conv_state + if conv_state is not None + else jnp.zeros( + (batch, self.config.gdn_conv_kernel_dim - 1, qkv.shape[-1]), + dtype=cfg.dtype, + ) ) - elif self.mesh is not None: - logical_rules = get_logical_axis_rules() recurrent_state_arg = ( - recurrent_state + recurrent_state.astype(state_dtype) if recurrent_state is not None else jnp.zeros( (batch, self.num_v_heads, self.head_k_dim, self.head_v_dim), - dtype=cfg.dtype, - out_sharding=state_sharding, + dtype=state_dtype, ) ) - # LENGTH, not None. The sequence axis was hardcoded to replicated, so - # ici_context_parallelism could never shard the GDN sequence while still - # consuming the context axis from the mesh -- which is why raising ctx - # made memory worse instead of better. The scan handles a sharded - # sequence via the two-pass affine composition in kernels/attention/gdn_cp.py. - # Either context axis can carry the sequence. LENGTH maps to both in the - # logical rules, so this is correct whichever one is configured. - cp_axes = gdn_context_axes(cfg) - cp_len = LENGTH if cp_axes else None - qkv_pspec = logical_to_mesh_axes((KV_BATCH, cp_len, KV_HEAD, None), mesh=self.mesh, rules=logical_rules) - g_beta_pspec = logical_to_mesh_axes((KV_BATCH, cp_len, KV_HEAD), mesh=self.mesh, rules=logical_rules) - state_pspec = logical_to_mesh_axes((KV_BATCH, KV_HEAD, None, None), mesh=self.mesh, rules=logical_rules) - # Keep every shard_map input/output batch spec consistent when replication is required. - qkv_pspec = remove_incompatible_mesh_axes_from_partition_spec( - qkv_pspec, - query.shape, - self.mesh, - dims=(0,), - allow_remove_axes=True, + conv_bias_arg = self.conv1d.bias.value if getattr(self.conv1d, "bias", None) is not None else None + + if self.mesh is not None: + logical_rules = self.config.logical_axis_rules + qkv_pspec = logical_to_mesh_axes((KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules) + b_a_pspec = logical_to_mesh_axes((KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules) + conv_state_pspec = logical_to_mesh_axes((KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules) + recurrent_state_pspec = logical_to_mesh_axes((KV_BATCH, None, None, None), mesh=self.mesh, rules=logical_rules) + + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=( + qkv_pspec, + b_a_pspec, + b_a_pspec, + P(), + P(), + P(), + P(), + conv_state_pspec, + recurrent_state_pspec, + ), + out_specs=( + qkv_pspec, + (conv_state_pspec, recurrent_state_pspec), + ), + check_vma=False, + ) + def shard_mapped_gdn( + qkv_val, + b_val, + a_val, + cw_val, + cb_val, + alog_val, + dt_val, + cs_val, + rs_val, + ): + return gdn_decoupled_conv1d( + qkv=qkv_val, + b=b_val, + a=a_val, + conv_weight=cw_val, + conv_bias=cb_val, + a_log=alog_val, + dt_bias=dt_val, + conv_state=cs_val, + recurrent_state=rs_val, + num_k_heads=self.num_k_heads, + num_v_heads=self.num_v_heads, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + conv_kernel_size=self.config.gdn_conv_kernel_dim, + chunk_size=self.config.gdn_chunk_size, + use_qk_norm_in_gdn=self.config.use_qk_norm_in_gdn, + compute_dtype=state_dtype, + ) + + gdn_step_fn = shard_mapped_gdn + core_attn_out, (next_conv_state, next_recurrent_state) = gdn_step_fn( + qkv, + b, + a, + self.conv1d.kernel.value, + conv_bias_arg, + self.A_log[...], + self.dt_bias[...], + conv_state_arg, + recurrent_state_arg, + ) + else: + gdn_step_fn = gdn_decoupled_conv1d + core_attn_out, (next_conv_state, next_recurrent_state) = gdn_step_fn( + qkv=qkv, + b=b, + a=a, + conv_weight=self.conv1d.kernel.value, + conv_bias=conv_bias_arg, + a_log=self.A_log[...], + dt_bias=self.dt_bias[...], + conv_state=conv_state_arg, + recurrent_state=recurrent_state_arg, + num_k_heads=self.num_k_heads, + num_v_heads=self.num_v_heads, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + conv_kernel_size=self.config.gdn_conv_kernel_dim, + chunk_size=self.config.gdn_chunk_size, + use_qk_norm_in_gdn=self.config.use_qk_norm_in_gdn, + compute_dtype=state_dtype, + ) + else: + # Perform the convolution. + conv_out = self.conv1d(conv_input, out_sharding=flat_sharding) + # Slice the output to match the original input sequence length. + conv_out = conv_out[:, -seq_len:, :] + qkv_conv = jax.nn.silu(conv_out.astype(jnp.float32)).astype(cfg.dtype) + # q_conv shape: (B, S, key_dim), k_conv shape: (B, S, key_dim), v_conv shape: (B, S, value_dim) + q_conv, k_conv, v_conv = jnp.split(qkv_conv, [self.key_dim, 2 * self.key_dim], axis=-1) + + # Reshape for multi-head processing + # query shape: (B, S, H_k, D_k) + query = jnp.reshape( + q_conv, + (batch, seq_len, self.num_k_heads, self.head_k_dim), + out_sharding=head_sharding, ) - g_beta_pspec = remove_incompatible_mesh_axes_from_partition_spec( - g_beta_pspec, - g.shape, - self.mesh, - dims=(0,), - allow_remove_axes=True, + # key shape: (B, S, H_k, D_k) + key = jnp.reshape( + k_conv, + (batch, seq_len, self.num_k_heads, self.head_k_dim), + out_sharding=head_sharding, ) - state_pspec = remove_incompatible_mesh_axes_from_partition_spec( - state_pspec, - recurrent_state_arg.shape, - self.mesh, - dims=(0,), - allow_remove_axes=True, + # value shape: (B, S, H_v, D_v) + value = jnp.reshape( + v_conv, + (batch, seq_len, self.num_v_heads, self.head_v_dim), + out_sharding=head_sharding, ) + # ========================================================================= + # STEP C: Gated Delta Rule Recurrence + # ========================================================================= + A_log = jnp.asarray(self.A_log[...], dtype=decay_dtype) + dt_bias = jnp.asarray(self.dt_bias[...], dtype=decay_dtype) if cfg.shard_mode == ShardMode.EXPLICIT: - # shard_map manualises the mesh axes it is given and will not insert a reshard - # for an operand whose layout differs from `in_specs`, so hand it arrays that - # already match. - query = jax.sharding.reshard(query, qkv_pspec) - key = jax.sharding.reshard(key, qkv_pspec) - value = jax.sharding.reshard(value, qkv_pspec) - g = jax.sharding.reshard(g, g_beta_pspec) - beta = jax.sharding.reshard(beta, g_beta_pspec) - recurrent_state_arg = jax.sharding.reshard(recurrent_state_arg, state_pspec) - - @functools.partial( - jax.shard_map, - mesh=self.mesh, - in_specs=( - qkv_pspec, # query - qkv_pspec, # key - qkv_pspec, # value - g_beta_pspec, # g - g_beta_pspec, # beta - state_pspec, # initial_state - ), - out_specs=( - qkv_pspec, # core_attn_out - state_pspec, # final_state - ), - check_vma=False, - ) - def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): - return jax_chunk_gated_delta_rule( - query=q, - key=k, - value=v, - g=g_val, - beta=beta_val, - chunk_size=cfg.gdn_chunk_size, - initial_state=init_h, + # Both are stored replicated but broadcast against (B, S, H_v) activations whose + # head axis is sharded, and explicit sharding requires broadcast operands to + # agree -- the same fix `_align_scale_with_normalized_axis` applies to the norm + # scales. + head_spec = jax.sharding.PartitionSpec(jax.typeof(a).sharding.spec[-1]) + A_log = jax.sharding.reshard(A_log, head_spec) + dt_bias = jax.sharding.reshard(dt_bias, head_spec) + # beta shape: (B, S, H_v) + beta = jax.nn.sigmoid(b) + # g shape: (B, S, H_v) + g = -jnp.exp(A_log) * jax.nn.softplus(a + dt_bias) + + if decoder_segment_ids is not None: + mask = decoder_segment_ids != 0 + # Apply mask by broadcasting to respective shapes + key = jnp.where(mask[..., None, None], key, 0.0) + value = jnp.where(mask[..., None, None], value, 0.0) + g = jnp.where(mask[..., None], g, 0.0) + + if self.num_v_heads > self.num_k_heads and self.num_v_heads % self.num_k_heads == 0: + repeats = self.num_v_heads // self.num_k_heads + # query shape after repeat: (B, S, H_v, D_k) + query = jnp.repeat(query, repeats, axis=2, out_sharding=head_sharding) + # key shape after repeat: (B, S, H_v, D_k) + key = jnp.repeat(key, repeats, axis=2, out_sharding=head_sharding) + + if seq_len == 1 and model_mode == MODEL_MODE_AUTOREGRESSIVE: + core_attn_out, next_recurrent_state = jax_ar_gated_delta_rule( + query, + key, + value, + g, + beta, + initial_state=recurrent_state, # pyrefly: ignore[bad-argument-type] use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, compute_dtype=cfg.dtype, - cp_axis=cp_axes or None, + ) + elif self.mesh is not None: + logical_rules = get_logical_axis_rules() + recurrent_state_arg = ( + recurrent_state + if recurrent_state is not None + else jnp.zeros( + (batch, self.num_v_heads, self.head_k_dim, self.head_v_dim), + dtype=state_dtype, + out_sharding=state_sharding, + ) + ) + # LENGTH, not None. The sequence axis was hardcoded to replicated, so + # ici_context_parallelism could never shard the GDN sequence while still + # consuming the context axis from the mesh -- which is why raising ctx + # made memory worse instead of better. The scan handles a sharded + # sequence via the two-pass affine composition in kernels/attention/gdn_cp.py. + # Either context axis can carry the sequence. LENGTH maps to both in the + # logical rules, so this is correct whichever one is configured. + cp_axes = gdn_context_axes(cfg) + cp_len = LENGTH if cp_axes else None + qkv_pspec = logical_to_mesh_axes((KV_BATCH, cp_len, KV_HEAD, None), mesh=self.mesh, rules=logical_rules) + g_beta_pspec = logical_to_mesh_axes((KV_BATCH, cp_len, KV_HEAD), mesh=self.mesh, rules=logical_rules) + state_pspec = logical_to_mesh_axes((KV_BATCH, KV_HEAD, None, None), mesh=self.mesh, rules=logical_rules) + # Keep every shard_map input/output batch spec consistent when replication is required. + qkv_pspec = remove_incompatible_mesh_axes_from_partition_spec( + qkv_pspec, + query.shape, + self.mesh, + dims=(0,), + allow_remove_axes=True, + ) + g_beta_pspec = remove_incompatible_mesh_axes_from_partition_spec( + g_beta_pspec, + g.shape, + self.mesh, + dims=(0,), + allow_remove_axes=True, + ) + state_pspec = remove_incompatible_mesh_axes_from_partition_spec( + state_pspec, + recurrent_state_arg.shape, + self.mesh, + dims=(0,), + allow_remove_axes=True, ) - core_attn_out, next_recurrent_state = shard_mapped_delta_rule(query, key, value, g, beta, recurrent_state_arg) - else: - core_attn_out, next_recurrent_state = jax_chunk_gated_delta_rule( - query, - key, - value, - g, - beta, - chunk_size=cfg.gdn_chunk_size, - initial_state=recurrent_state, - use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, - compute_dtype=cfg.dtype, - ) + if cfg.shard_mode == ShardMode.EXPLICIT: + # shard_map manualises the mesh axes it is given and will not insert a reshard + # for an operand whose layout differs from `in_specs`, so hand it arrays that + # already match. + query = jax.sharding.reshard(query, qkv_pspec) + key = jax.sharding.reshard(key, qkv_pspec) + value = jax.sharding.reshard(value, qkv_pspec) + g = jax.sharding.reshard(g, g_beta_pspec) + beta = jax.sharding.reshard(beta, g_beta_pspec) + recurrent_state_arg = jax.sharding.reshard(recurrent_state_arg, state_pspec) + + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=( + qkv_pspec, # query + qkv_pspec, # key + qkv_pspec, # value + g_beta_pspec, # g + g_beta_pspec, # beta + state_pspec, # initial_state + ), + out_specs=( + qkv_pspec, # core_attn_out + state_pspec, # final_state + ), + check_vma=False, + ) + def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h): + return jax_chunk_gated_delta_rule( + query=q, + key=k, + value=v, + g=g_val, + beta=beta_val, + chunk_size=cfg.gdn_chunk_size, + initial_state=init_h, + use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, + compute_dtype=cfg.dtype, + cp_axis=cp_axes or None, + ) + + core_attn_out, next_recurrent_state = shard_mapped_delta_rule(query, key, value, g, beta, recurrent_state_arg) + else: + core_attn_out, next_recurrent_state = jax_chunk_gated_delta_rule( + query, + key, + value, + g, + beta, + chunk_size=cfg.gdn_chunk_size, + initial_state=recurrent_state, + use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn, + compute_dtype=cfg.dtype, + ) if model_mode != MODEL_MODE_TRAIN and active_cache is not None: assert next_conv_state is not None diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py index d870026e3..16b2a1ce1 100644 --- a/tests/unit/pyconfig_test.py +++ b/tests/unit/pyconfig_test.py @@ -45,6 +45,7 @@ def test_gmm_v2_heuristic_tiling_requires_gmm_v2(self): with self.assertRaisesRegex(ValueError, "`use_gmm_v2_heuristic_tiling=True` requires `use_gmm_v2=True`."): pyconfig.initialize( [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()], + skip_jax_distributed_system=True, use_gmm_v2_heuristic_tiling=True, use_gmm_v2=False, ) @@ -61,6 +62,7 @@ def test_gdn_context_parallelism_rejects_load_balance(self): model_name="qwen3-next-80b-a3b", ici_context_parallelism=4, context_parallel_load_balance=True, + skip_jax_distributed_system=True, ) def test_gdn_context_parallelism_accepts_load_balance_off(self):