Skip to content

Commit 277bd70

Browse files
changlanchanglan
committed
[Qwen] Use fsdp for B200
GitOrigin-RevId: 446381b
1 parent 9b6d5f3 commit 277bd70

3 files changed

Lines changed: 33 additions & 74 deletions

File tree

‎axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/qwen3-30B-A3B-bs16m-seq32k.txt‎

Lines changed: 5 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -175,37 +175,11 @@ mesh_rules[1][1].config_modifiers[0].mesh_shape['dcn_mesh_shape'][4]: 1
175175
mesh_rules[1][1].config_modifiers[0].mesh_shape['dcn_mesh_shape'][5]: 1
176176
mesh_rules[1][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
177177
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
178-
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy']: 'jax._src.ad_checkpoint.dots_with_no_batch_dims_saveable'
179-
mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.ReplaceLayerConfigModifier'
180-
mesh_rules[1][1].config_modifiers[2].source_config.activation: 'nn.relu'
181-
mesh_rules[1][1].config_modifiers[2].source_config.corr_loss_weight: 0
182-
mesh_rules[1][1].config_modifiers[2].source_config.dim_to_mesh_axis_map['me']: PartitionSpec(None, None)
183-
mesh_rules[1][1].config_modifiers[2].source_config.dim_to_mesh_axis_map['emh']: PartitionSpec('expert', ('fsdp', 'seq'), 'model')
184-
mesh_rules[1][1].config_modifiers[2].source_config.dim_to_mesh_axis_map['ehm']: PartitionSpec('expert', 'model', ('fsdp', 'seq'))
185-
mesh_rules[1][1].config_modifiers[2].source_config.dropout.klass: 'axlearn.common.layers.Dropout'
186-
mesh_rules[1][1].config_modifiers[2].source_config.expert_parallel_capacity: 1.25
187-
mesh_rules[1][1].config_modifiers[2].source_config.gating.eval_capacity_factor: 2.0
188-
mesh_rules[1][1].config_modifiers[2].source_config.gating.gating_logit_cap: 0.0
189-
mesh_rules[1][1].config_modifiers[2].source_config.gating.klass: 'axlearn.common.mixture_of_experts.Top2Gating'
190-
mesh_rules[1][1].config_modifiers[2].source_config.gating.mask_dtype: 'jax.numpy.int32'
191-
mesh_rules[1][1].config_modifiers[2].source_config.gating.train_capacity_factor: 2.0
192-
mesh_rules[1][1].config_modifiers[2].source_config.input_dim_to_partition_spec['bsm']: PartitionSpec(('expert', 'data', 'fsdp'), 'seq', None)
193-
mesh_rules[1][1].config_modifiers[2].source_config.interpret: False
194-
mesh_rules[1][1].config_modifiers[2].source_config.klass: 'axlearn.common.mixture_of_experts.ApproximateTokenDropFreeMoE'
195-
mesh_rules[1][1].config_modifiers[2].source_config.load_balance_loss_weight: 0.01
196-
mesh_rules[1][1].config_modifiers[2].source_config.norm.eps: 1e-08
197-
mesh_rules[1][1].config_modifiers[2].source_config.norm.forward_dtype: 'jax.numpy.float32'
198-
mesh_rules[1][1].config_modifiers[2].source_config.norm.klass: 'axlearn.common.layers.LayerNorm'
199-
mesh_rules[1][1].config_modifiers[2].source_config.outer_batch: 1
200-
mesh_rules[1][1].config_modifiers[2].source_config.output_dim_to_partition_spec['bsm']: PartitionSpec(('expert', 'data', 'fsdp'), 'seq', 'model')
201-
mesh_rules[1][1].config_modifiers[2].source_config.output_dim_to_partition_spec['emh']: PartitionSpec('expert', None, 'model')
202-
mesh_rules[1][1].config_modifiers[2].source_config.output_dim_to_partition_spec['ehm']: PartitionSpec('expert', 'model', None)
203-
mesh_rules[1][1].config_modifiers[2].source_config.residual_weight: 1.0
204-
mesh_rules[1][1].config_modifiers[2].source_config.router_z_loss_weight: 0.0
205-
mesh_rules[1][1].config_modifiers[2].source_config.stochastic_depth.klass: 'axlearn.common.layers.StochasticDepth'
206-
mesh_rules[1][1].config_modifiers[2].source_config.stochastic_depth.mode: 'row'
207-
mesh_rules[1][1].config_modifiers[2].source_config.structure: 'prenorm'
208-
mesh_rules[1][1].config_modifiers[2].target_cls: 'axlearn.common.mixture_of_experts.TransformerFeedForwardDropFreeMoE'
178+
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
179+
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: ''
180+
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*flash_residuals|.*[kqv]_proj|.*o_proj|.*linear1_[01]|.*linear2(_[01])?|.*gate_weights|.*gate_assignment'
181+
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
182+
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
209183
mesh_rules[1][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
210184
mesh_shape[0]: 1
211185
mesh_shape[1]: -1

‎axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/qwen3-30B-A3B-bs8m-seq8k.txt‎

Lines changed: 5 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -175,37 +175,11 @@ mesh_rules[1][1].config_modifiers[0].mesh_shape['dcn_mesh_shape'][4]: 1
175175
mesh_rules[1][1].config_modifiers[0].mesh_shape['dcn_mesh_shape'][5]: 1
176176
mesh_rules[1][1].config_modifiers[1].klass: 'axlearn.common.trainer_config_modifier.RematSpecModifier'
177177
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['prevent_cse']: False
178-
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy']: 'jax._src.ad_checkpoint.dots_with_no_batch_dims_saveable'
179-
mesh_rules[1][1].config_modifiers[2].klass: 'axlearn.common.trainer_config_modifier.ReplaceLayerConfigModifier'
180-
mesh_rules[1][1].config_modifiers[2].source_config.activation: 'nn.relu'
181-
mesh_rules[1][1].config_modifiers[2].source_config.corr_loss_weight: 0
182-
mesh_rules[1][1].config_modifiers[2].source_config.dim_to_mesh_axis_map['me']: PartitionSpec(None, None)
183-
mesh_rules[1][1].config_modifiers[2].source_config.dim_to_mesh_axis_map['emh']: PartitionSpec('expert', ('fsdp', 'seq'), 'model')
184-
mesh_rules[1][1].config_modifiers[2].source_config.dim_to_mesh_axis_map['ehm']: PartitionSpec('expert', 'model', ('fsdp', 'seq'))
185-
mesh_rules[1][1].config_modifiers[2].source_config.dropout.klass: 'axlearn.common.layers.Dropout'
186-
mesh_rules[1][1].config_modifiers[2].source_config.expert_parallel_capacity: 1.25
187-
mesh_rules[1][1].config_modifiers[2].source_config.gating.eval_capacity_factor: 2.0
188-
mesh_rules[1][1].config_modifiers[2].source_config.gating.gating_logit_cap: 0.0
189-
mesh_rules[1][1].config_modifiers[2].source_config.gating.klass: 'axlearn.common.mixture_of_experts.Top2Gating'
190-
mesh_rules[1][1].config_modifiers[2].source_config.gating.mask_dtype: 'jax.numpy.int32'
191-
mesh_rules[1][1].config_modifiers[2].source_config.gating.train_capacity_factor: 2.0
192-
mesh_rules[1][1].config_modifiers[2].source_config.input_dim_to_partition_spec['bsm']: PartitionSpec(('expert', 'data', 'fsdp'), 'seq', None)
193-
mesh_rules[1][1].config_modifiers[2].source_config.interpret: False
194-
mesh_rules[1][1].config_modifiers[2].source_config.klass: 'axlearn.common.mixture_of_experts.ApproximateTokenDropFreeMoE'
195-
mesh_rules[1][1].config_modifiers[2].source_config.load_balance_loss_weight: 0.01
196-
mesh_rules[1][1].config_modifiers[2].source_config.norm.eps: 1e-08
197-
mesh_rules[1][1].config_modifiers[2].source_config.norm.forward_dtype: 'jax.numpy.float32'
198-
mesh_rules[1][1].config_modifiers[2].source_config.norm.klass: 'axlearn.common.layers.LayerNorm'
199-
mesh_rules[1][1].config_modifiers[2].source_config.outer_batch: 1
200-
mesh_rules[1][1].config_modifiers[2].source_config.output_dim_to_partition_spec['bsm']: PartitionSpec(('expert', 'data', 'fsdp'), 'seq', 'model')
201-
mesh_rules[1][1].config_modifiers[2].source_config.output_dim_to_partition_spec['emh']: PartitionSpec('expert', None, 'model')
202-
mesh_rules[1][1].config_modifiers[2].source_config.output_dim_to_partition_spec['ehm']: PartitionSpec('expert', 'model', None)
203-
mesh_rules[1][1].config_modifiers[2].source_config.residual_weight: 1.0
204-
mesh_rules[1][1].config_modifiers[2].source_config.router_z_loss_weight: 0.0
205-
mesh_rules[1][1].config_modifiers[2].source_config.stochastic_depth.klass: 'axlearn.common.layers.StochasticDepth'
206-
mesh_rules[1][1].config_modifiers[2].source_config.stochastic_depth.mode: 'row'
207-
mesh_rules[1][1].config_modifiers[2].source_config.structure: 'prenorm'
208-
mesh_rules[1][1].config_modifiers[2].target_cls: 'axlearn.common.mixture_of_experts.TransformerFeedForwardDropFreeMoE'
178+
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].fn: 'axlearn.common.utils.save_and_offload_only_these_names_regex'
179+
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_offloaded: ''
180+
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].names_which_can_be_saved: '.*flash_residuals|.*[kqv]_proj|.*o_proj|.*linear1_[01]|.*linear2(_[01])?|.*gate_weights|.*gate_assignment'
181+
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_dst: 'pinned_host'
182+
mesh_rules[1][1].config_modifiers[1].remat_policies['model.decoder.transformer.layer']['policy'].offload_src: 'device'
209183
mesh_rules[1][1].klass: 'axlearn.common.trainer_config_modifier.ChainConfigModifier'
210184
mesh_shape[0]: 1
211185
mesh_shape[1]: -1

‎axlearn/experiments/text/gpt/qwen.py‎

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
from typing import Any, Literal, Optional, Sequence, Union
88

99
import jax.numpy as jnp
10-
from jax.ad_checkpoint import checkpoint_policies as jax_remat_policies
1110

1211
from axlearn.common import causal_lm, config, decoder
1312
from axlearn.common.attention import (
@@ -24,16 +23,11 @@
2423
from axlearn.common.flash_attention.layer import FlashAttention
2524
from axlearn.common.flash_attention.remat import save_or_offload_flash_attention_policy
2625
from axlearn.common.layers import RMSNorm
27-
from axlearn.common.mixture_of_experts import (
28-
ApproximateTokenDropFreeMoE,
29-
TopKDropFreeGating,
30-
TransformerFeedForwardDropFreeMoE,
31-
)
26+
from axlearn.common.mixture_of_experts import TopKDropFreeGating, TransformerFeedForwardDropFreeMoE
3227
from axlearn.common.trainer_config_modifier import (
3328
ChainConfigModifier,
3429
MeshShapeModifier,
3530
RematSpecModifier,
36-
ReplaceLayerConfigModifier,
3731
)
3832
from axlearn.common.utils import (
3933
HybridMeshShape,
@@ -177,14 +171,31 @@ def get_trainer_kwargs(
177171
remat_policies={
178172
"model.decoder.transformer.layer": RematSpec(
179173
prevent_cse=False,
180-
policy=jax_remat_policies.dots_with_no_batch_dims_saveable,
174+
policy=config_for_function(
175+
save_and_offload_only_these_names_regex
176+
).set(
177+
names_which_can_be_saved="|".join(
178+
[
179+
RematRegexSavePatterns.FLASH_CONTEXT.value,
180+
RematRegexSavePatterns.QKV_PROJ.value,
181+
RematRegexSavePatterns.O_PROJ.value,
182+
RematRegexSavePatterns.LINEAR1_X.value,
183+
RematRegexSavePatterns.LINEAR2_X.value,
184+
RematRegexSavePatterns.MOE_GATING.value,
185+
]
186+
),
187+
names_which_can_be_offloaded="|".join([]),
188+
offload_src="device",
189+
offload_dst="pinned_host",
190+
),
181191
),
182192
}
183193
),
184-
ReplaceLayerConfigModifier.default_config().set(
185-
target_cls=TransformerFeedForwardDropFreeMoE,
186-
source_config=ApproximateTokenDropFreeMoE.default_config(),
187-
),
194+
# TODO: Fix the ApproximateTokenDropFreeMoE performance issue on GPUs.
195+
# ReplaceLayerConfigModifier.default_config().set(
196+
# target_cls=TransformerFeedForwardDropFreeMoE,
197+
# source_config=ApproximateTokenDropFreeMoE.default_config(),
198+
# ),
188199
],
189200
),
190201
),

0 commit comments

Comments
 (0)