diff --git a/packages/paddlefleet_ops/third_party/sonic-moe b/packages/paddlefleet_ops/third_party/sonic-moe index 226ebb059..cd918c5f0 160000 --- a/packages/paddlefleet_ops/third_party/sonic-moe +++ b/packages/paddlefleet_ops/third_party/sonic-moe @@ -1 +1 @@ -Subproject commit 226ebb059e373540731167b27d152ae329dcbfb2 +Subproject commit cd918c5f0bf9ac8e251b3f30223f060cef8df600 diff --git a/src/paddlefleet/transformer/moe/fusion_layer_utils.py b/src/paddlefleet/transformer/moe/fusion_layer_utils.py index 7d7106293..7ae1f0d9c 100644 --- a/src/paddlefleet/transformer/moe/fusion_layer_utils.py +++ b/src/paddlefleet/transformer/moe/fusion_layer_utils.py @@ -32,21 +32,29 @@ tokens_zip_unique_add_with_subbatch, ) +_scatter_router_scores_i32 = None if paddlefleet_ops.is_sonic_moe_available(): from paddlefleet_ops.sonicmoe.enums import ActivationType from paddlefleet_ops.sonicmoe.ernie_compat.deepep_metadata import ( deepep_topk_to_sonic_metadata, + deepep_topk_to_sonic_metadata_with_scales, ) from paddlefleet_ops.sonicmoe.ernie_compat.mlp_node_v2 import ( _differentiable_router_scores, ) from paddlefleet_ops.sonicmoe.functional import ( _DownProjection, - _refresh_fp8_config, _UpProjection, ) from paddlefleet_ops.sonicmoe.functional.utils import enable_fp8 + try: + from paddlefleet_ops.sonicmoe.quack_utils.blockscaled_fp8_gemm import ( + _scatter_router_scores_i32, + ) + except (ImportError, RuntimeError): + pass + logger = logging.getLogger(__name__) @@ -82,6 +90,57 @@ def _make_sonic_fp8_weight_carrier(weight): return _SonicFP8WeightCarrier.apply(weight) +class _SonicRouterScoresFromMetadata(paddle.autograd.PyLayer): + @staticmethod + def forward(ctx, topk_scores, metadata_scores, score_src_idx): + if len(topk_scores.shape) != 2: + raise ValueError( + f"topk_scores: expected rank 2, got shape {topk_scores.shape}" + ) + if len(metadata_scores.shape) != 1: + raise ValueError( + "metadata_scores: expected rank 1, got shape " + f"{metadata_scores.shape}" + ) + if len(score_src_idx.shape) != 1: + raise ValueError( + f"score_src_idx: expected rank 1, got shape {score_src_idx.shape}" + ) + if metadata_scores.shape[0] < score_src_idx.shape[0]: + raise ValueError( + "metadata_scores must include every real score referenced by " + f"score_src_idx; got {metadata_scores.shape[0]} scores and " + f"{score_src_idx.shape[0]} indices" + ) + if "int32" not in str(score_src_idx.dtype): + raise ValueError( + f"score_src_idx: expected int32, got {score_src_idx.dtype}" + ) + metadata_scores.stop_gradient = True + score_src_idx.stop_gradient = True + ctx.save_for_backward(score_src_idx) + ctx.input_shape = list(topk_scores.shape) + ctx.n_total = int(topk_scores.shape[0]) * int(topk_scores.shape[1]) + scores = metadata_scores.clone() + scores.stop_gradient = topk_scores.stop_gradient + return scores + + @staticmethod + def backward(ctx, grad_out): + (score_src_idx,) = ctx.saved_tensor() + if _scatter_router_scores_i32 is None: + raise RuntimeError( + "SonicMoE metadata router score backward requires " + "paddlefleet_ops.sonicmoe.quack_utils.blockscaled_fp8_gemm." + "_scatter_router_scores_i32; update paddlefleet_ops or use " + "the differentiable router-score fallback." + ) + grad_flat = _scatter_router_scores_i32( + grad_out.contiguous(), score_src_idx, ctx.n_total + ) + return grad_flat.reshape(ctx.input_shape), None, None + + class UnZipNode: """ UnZipNode 类用于对输入的token 矩阵根据分发索引进行解压操作,得到专家需要处理的 token。 @@ -3187,6 +3246,7 @@ def run_sonic_moe( tokens_per_expert=None, fp8_scale=None, fp8_combine_grad_handle=None, + fp8_config=None, ): T = hidden_states.shape[0] stream_id = paddle.device.current_stream() @@ -3198,64 +3258,88 @@ def run_sonic_moe( paddle.int32 ) - # paddle.cast is not a no-op for matching dtype (it allocates + copies), - # so cast once with a guard instead of twice below. Allgather feeds int32 - # (zero copies); deepep feeds int64 (one copy instead of two). - topk_indices_i32 = ( - topk_indices - if topk_indices.dtype == paddle.int32 - else topk_indices.cast(paddle.int32) - ) - - ( - expert_frequency_offset, - x_gather_idx, - s_scatter_idx, - s_reverse_scatter_idx, - num_activated_expert_per_token_offset, - _router_scores, - TK_padded, - total_pad_rows, - _N_recv, - _score_src_idx, - ) = deepep_topk_to_sonic_metadata( - topk_indices_i32, - topk_scores, - tokens_per_expert, - E, - block=128 if fp8 else 1, - ) + fp8_scale_packed = None + if fp8 and fp8_scale is not None: + ( + expert_frequency_offset, + x_gather_idx, + s_scatter_idx, + s_reverse_scatter_idx, + num_activated_expert_per_token_offset, + _router_scores, + TK_padded, + total_pad_rows, + _N_recv, + _score_src_idx, + fp8_scale_packed, + ) = deepep_topk_to_sonic_metadata_with_scales( + topk_indices.cast(paddle.int32), + topk_scores, + tokens_per_expert, + E, + fp8_scale, + int(hidden_states.shape[1]), + block=128, + ) + else: + ( + expert_frequency_offset, + x_gather_idx, + s_scatter_idx, + s_reverse_scatter_idx, + num_activated_expert_per_token_offset, + _router_scores, + TK_padded, + total_pad_rows, + _N_recv, + _score_src_idx, + ) = deepep_topk_to_sonic_metadata( + topk_indices.cast(paddle.int32), + topk_scores, + tokens_per_expert, + E, + block=128 if fp8 else 1, + ) s_scatter_idx.stop_gradient = True activation_type = ActivationType("swiglu") total_expert_freq = TK_padded - scores_for_down = _differentiable_router_scores( - topk_scores, - topk_indices_i32, - num_activated_expert_per_token_offset, - TK_padded - total_pad_rows, - TK_padded, - E, - score_src_idx=_score_src_idx, - ) + if _score_src_idx is not None and _scatter_router_scores_i32 is not None: + scores_for_down = _SonicRouterScoresFromMetadata.apply( + topk_scores, _router_scores, _score_src_idx + ) + else: + scores_for_down = _differentiable_router_scores( + topk_scores, + topk_indices.cast(paddle.int32), + num_activated_expert_per_token_offset, + TK_padded - total_pad_rows, + TK_padded, + E, + score_src_idx=_score_src_idx, + ) fp8_hidden_states = None if fp8_scale is not None: - fp8_hidden_states = (hidden_states, fp8_scale) + fp8_hidden_states = ( + (hidden_states, fp8_scale, fp8_scale_packed) + if fp8_scale_packed is not None + else (hidden_states, fp8_scale) + ) - if fp8: - w1_sonic = _make_sonic_fp8_weight_carrier(w1) - w2_sonic = _make_sonic_fp8_weight_carrier(w2) - else: - w1_sonic = w1.permute([1, 2, 0]) - w2_sonic = w2.permute([1, 2, 0]) + # if fp8: + # w1_sonic = _make_sonic_fp8_weight_carrier(w1) + # w2_sonic = _make_sonic_fp8_weight_carrier(w2) + # else: + # w1_sonic = w1.permute([1, 2, 0]) + # w2_sonic = w2.permute([1, 2, 0]) with enable_fp8(fp8): - _refresh_fp8_config() + # _refresh_fp8_config() y1, z = _UpProjection.apply( hidden_states, - w1_sonic, + w1, None, expert_frequency_offset, total_expert_freq, @@ -3270,11 +3354,12 @@ def run_sonic_moe( is_inference_mode_enabled=False, use_low_precision_postact_buffer=False, prequant_activation_payload=fp8_hidden_states, + fp8_config=fp8_config, ) hidden_states = _DownProjection.apply( y1, z, - w2_sonic, + w2, None, scores_for_down, s_scatter_idx, @@ -3290,6 +3375,7 @@ def run_sonic_moe( activation_type, None, fp8_combine_grad_handle, + fp8_config=fp8_config, ) return hidden_states diff --git a/src/paddlefleet/transformer/moe/moe_expert.py b/src/paddlefleet/transformer/moe/moe_expert.py index d31d00451..4758c8138 100644 --- a/src/paddlefleet/transformer/moe/moe_expert.py +++ b/src/paddlefleet/transformer/moe/moe_expert.py @@ -40,12 +40,27 @@ try: from paddlefleet_ops import deep_gemm as paddlefleet_deep_gemm - from paddlefleet_ops.sonicmoe.functional import clear_all_fp8_weight_caches + from paddlefleet_ops.sonicmoe.functional import ( + _refresh_fp8_config, + clear_all_fp8_weight_caches, + ) from paddlefleet_ops.sonicmoe.quack_utils import quantize_native_fp8_weights except (ImportError, RuntimeError): pass +try: + from paddlefleet_ops.sonicmoe.ernie_compat.weight_layout_fusion import ( + fused_grouped_w1_to_sonic, + fused_sonic_w1_to_grouped, + fused_transpose_w2_layout, + ) +except (ImportError, RuntimeError): + fused_grouped_w1_to_sonic = None + fused_sonic_w1_to_grouped = None + fused_transpose_w2_layout = None + + class BMMFunction(paddle.autograd.PyLayer): @staticmethod def forward(ctx, x, y, batch_sizes, trans_y=False): @@ -459,25 +474,34 @@ def _is_tensor_initialized(tensor): @staticmethod def _grouped_w1_to_sonic(weight): - target_shape = [weight.shape[0], weight.shape[2], weight.shape[1]] - gate, up = paddle.chunk(weight, 2, axis=-1) - gate = gate.transpose([0, 2, 1]) - up = up.transpose([0, 2, 1]) - return paddle.stack([gate, up], axis=2).reshape(target_shape) + if fused_grouped_w1_to_sonic is not None: + return fused_grouped_w1_to_sonic(weight) + else: + target_shape = [weight.shape[0], weight.shape[2], weight.shape[1]] + gate, up = paddle.chunk(weight, 2, axis=-1) + gate = gate.transpose([0, 2, 1]) + up = up.transpose([0, 2, 1]) + return paddle.stack([gate, up], axis=2).reshape(target_shape) @staticmethod def _sonic_w1_to_grouped(weight): - target_shape = [weight.shape[0], weight.shape[2], weight.shape[1]] - weight = weight.reshape([weight.shape[0], -1, 2, weight.shape[2]]) - gate = weight[:, :, 0, :].transpose([0, 2, 1]) - up = weight[:, :, 1, :].transpose([0, 2, 1]) - return paddle.concat([gate, up], axis=-1) + if fused_sonic_w1_to_grouped is not None: + return fused_sonic_w1_to_grouped(weight) + else: + target_shape = [weight.shape[0], weight.shape[2], weight.shape[1]] + weight = weight.reshape([weight.shape[0], -1, 2, weight.shape[2]]) + gate = weight[:, :, 0, :].transpose([0, 2, 1]) + up = weight[:, :, 1, :].transpose([0, 2, 1]) + return paddle.concat([gate, up], axis=-1) @staticmethod def _transpose_w2_layout(weight): # if not SonicMoEExpert._is_tensor_initialized(weight): # return weight - return weight.transpose([0, 2, 1]) + if fused_transpose_w2_layout is not None: + return fused_transpose_w2_layout(weight) + else: + return weight.transpose([0, 2, 1]) @staticmethod def _assign_tensor(tensor, value): @@ -511,10 +535,15 @@ def __init__( self.hidden_size = self.config.hidden_size self.K = topk self._weights_layout = self._GROUPED_LAYOUT + self.sonic_moe_config = _refresh_fp8_config() + self.sonic_moe_config.enabled = self.config.fp8 is not None + self.sonic_moe_config.fp8_wgrad = self.config.fp8_wgrad + self.sonic_moe_config.fuse_y1_quant = True + self.sonic_moe_config.fuse_y1_bf16_trunc = True - def _convert_grad_layout(self, param, converter): + def _convert_grad_layout(self, param, converter, convert_main_grad=True): main_grad = getattr(param, "main_grad", None) - if main_grad is not None: + if convert_main_grad and main_grad is not None: self._assign_tensor(main_grad, converter(main_grad)) if param.grad is not None and ( main_grad is None or param.grad.data_ptr() != main_grad.data_ptr() @@ -536,7 +565,15 @@ def _convert_layout( param.get_tensor()._set_dims([shape[0], shape[2], shape[1]]) else: self._assign_tensor(param, converter(param)) - self._convert_grad_layout(param, converter) + # weight2's main_grad stays in the original grouped layout + # ([E, I, H]); its grouped->sonic->grouped transpose is elided and + # the down-proj bf16/fp8 wgrad accumulates into it via a permute + # view. weight1's main_grad must still be converted because the + # grouped->sonic conversion also interleaves gate/up (a perfect + # shuffle the wgrad kernel does not undo). + self._convert_grad_layout( + param, converter, convert_main_grad=(param is self.weight1) + ) self._weights_layout = target_layout def convert_weights_to_sonic_layout(self): @@ -564,15 +601,15 @@ def quant_weight(self): self.convert_weights_to_sonic_layout() payload = quantize_native_fp8_weights( - self.weight1.permute([1, 2, 0]), - self.weight2.permute([1, 2, 0]), + self.weight1, + self.weight2, ) assert payload["format"] == "1x32", ( f"quant strategy {payload.get('format')} is not supported." ) w1_fp8, w1_scale, w1t_fp8, w1t_scale = payload["w1"] w2_fp8, w2_scale, w2t_fp8, w2t_scale = payload["w2"] - self.weight1.fp8 = (w1_fp8.mT, w1_scale) + self.weight1.fp8 = (w1_fp8, w1_scale) self.weight1.transposed_fp8 = (w1t_fp8, w1t_scale) self.weight2.fp8 = (w2_fp8, w2_scale) self.weight2.transposed_fp8 = (w2t_fp8, w2t_scale) @@ -583,6 +620,14 @@ def clear_fp8_weights(self): weight.fp8 = None weight.transposed_fp8 = None + def need_quant_weight(self): + for w in [self.weight1, self.weight2]: + if not hasattr(w, "fp8") or w.fp8 is None: + return True + if not hasattr(w, "transposed_fp8") or w.transposed_fp8 is None: + return True + return False + def forward( self, hidden_states, @@ -594,6 +639,8 @@ def forward( fp8_combine_grad_handle=None, ): self.convert_weights_to_sonic_layout() + if self.sonic_moe_config.enabled is True and self.need_quant_weight(): + self.quant_weight() hidden_states = run_sonic_moe( hidden_states, topk_indices, @@ -606,6 +653,7 @@ def forward( tokens_per_expert=tokens_per_expert, fp8_scale=fp8_scale, fp8_combine_grad_handle=fp8_combine_grad_handle, + fp8_config=self.sonic_moe_config, ) return hidden_states diff --git a/src/paddlefleet/transformer/moe/moe_layer.py b/src/paddlefleet/transformer/moe/moe_layer.py index 2e8acd53c..68f1f5ec0 100644 --- a/src/paddlefleet/transformer/moe/moe_layer.py +++ b/src/paddlefleet/transformer/moe/moe_layer.py @@ -180,6 +180,9 @@ def __init__( self.using_sonic_moe = self.config.using_sonic_moe self.fp8_dispatch = bool(config.fp8) self.fp8_wgrad = config.fp8_wgrad + self.fp8_dispatch_bwd = ( + self.fp8_dispatch and self.using_sonic_moe and self.fp8_wgrad + ) self.moe_expert_fusion = config.moe_expert_fusion self.moe_subbatch_token_num_after_dispatch = ( config.moe_subbatch_token_num_after_dispatch @@ -749,7 +752,7 @@ def combine( hidden_states, combine_overlap_handle, use_rr_deepep_combine=self.use_rr_deepep_combine, - fp8_dispatch=self.fp8_dispatch and self.using_sonic_moe, + fp8_dispatch=self.fp8_dispatch_bwd, combine_grad_handle=fp8_combine_grad_handle, ) @@ -917,9 +920,8 @@ def fusion_moe_forward( self.num_experts_per_tok, tokens_per_expert, ) - fp8_combine_grad_handle = ( - {} if self.fp8_dispatch and self.using_sonic_moe else None - ) + fp8_combine_grad_handle = {} if self.fp8_dispatch_bwd else None + # fp8_combine_grad_handle = None with profile("fusion_mlp"): if self._use_hybrid_ep_fusion(): diff --git a/tests/multi_card_tests/moe/test_sonic_moe_ep.py b/tests/multi_card_tests/moe/test_sonic_moe_ep.py index 86bf687e4..25041ec05 100644 --- a/tests/multi_card_tests/moe/test_sonic_moe_ep.py +++ b/tests/multi_card_tests/moe/test_sonic_moe_ep.py @@ -264,11 +264,16 @@ def _build_moe_layer( transformer_config, num_experts=self.n_routed_experts, ) - return MoELayer( + moe_layer = MoELayer( transformer_config, transformer_layer_spec.sublayers_spec.mlp.extra_kwargs["sublayers"], pg_collection or self.pg_collection, ) + mix_precision_utils.MixPrecisionLayer(moe_layer, dtype="bfloat16") + for param in moe_layer.parameters(): + if hasattr(param, "main_grad") and param.main_grad is None: + param.main_grad = paddle.zeros_like(param, dtype=paddle.float32) + return moe_layer @staticmethod def _expert_slice_for_rank(tensor, ep_rank, ep_size): @@ -352,7 +357,7 @@ def _run_forward_backward(self, moe_layer, input_data): master_grad=True, master_weight=True, ) - mix_precision_utils.MixPrecisionLayer(moe_layer, dtype="bfloat16") + # mix_precision_utils.MixPrecisionLayer(moe_layer, dtype="bfloat16") hidden_states = input_data.detach().clone() hidden_states.stop_gradient = False with paddle.amp.auto_cast(level="O2", dtype="bfloat16"): @@ -605,9 +610,55 @@ def run_test_ep_precision(self): print("Expert-parallel FP8 precision checks passed!") + def run_test_bf16_wgrad(self): + acc_steps = 1 + + paddle.seed(self.seed) + model_parallel_cuda_manual_seed(self.seed) + moe_layer_sonic_bf16 = self._build_moe_layer(using_sonic_moe=True) + paddle.seed(self.seed) + model_parallel_cuda_manual_seed(self.seed) + moe_layer_sonic_fp8 = self._build_moe_layer( + using_sonic_moe=True, + fp8="e4m3", + fp8_wgrad=False, + ) + + input_data_list = [] + for step_idx in range(acc_steps): + paddle.seed(self.seed + step_idx) + input_data_list.append( + paddle.randn( + [4, 256, self.hidden_size], + dtype=paddle.bfloat16, + ) + ) + + output_bf16, grads_bf16 = self._run_accumulated_forward_backward( + moe_layer_sonic_bf16, input_data_list + ) + output_fp8, grads_fp8 = self._run_accumulated_forward_backward( + moe_layer_sonic_fp8, input_data_list + ) + + fp8_tol = 5e-3 + self._assert_tensor_diff_less( + output_fp8, + output_bf16, + tol=fp8_tol, + title="Sonic-MoE FP8 vs BF16 final output", + ) + self._assert_grad_diff_less( + grads_fp8, + grads_bf16, + tol=fp8_tol, + title="Sonic-MoE FP8 vs BF16 accumulated grad", + ) + def test_sonic_moe_all(self): self.run_test_sonic_moe_ep_grad_accumulation() self.run_test_ep_precision() + self.run_test_bf16_wgrad() if __name__ == "__main__": diff --git a/tests/single_card_tests/ai_edited_test/transformer/test_ai_moe_layer_7.py b/tests/single_card_tests/ai_edited_test/transformer/test_ai_moe_layer_7.py index 3dc365919..d46d5409e 100644 --- a/tests/single_card_tests/ai_edited_test/transformer/test_ai_moe_layer_7.py +++ b/tests/single_card_tests/ai_edited_test/transformer/test_ai_moe_layer_7.py @@ -286,6 +286,7 @@ def test_combine_delegates(self): layer.use_rr_deepep_combine = False layer.fp8_dispatch = False layer.using_sonic_moe = False + layer.fp8_dispatch_bwd = False hidden = paddle.randn([4, 64]) layer.token_dispatcher._comm_manager.combine.return_value = hidden result = layer.combine(hidden) diff --git a/tests/single_card_tests/model/test_gpt_model_sonic_moe.py b/tests/single_card_tests/model/test_gpt_model_sonic_moe.py index 71c29cf88..aa4b137f7 100644 --- a/tests/single_card_tests/model/test_gpt_model_sonic_moe.py +++ b/tests/single_card_tests/model/test_gpt_model_sonic_moe.py @@ -16,10 +16,12 @@ import subprocess import unittest +import numpy as np import paddle import paddle.nn.functional as F import paddlefleet_ops from paddle.distributed import fleet +from paddle.distributed.fleet.utils import mix_precision_utils from paddlefleet_ops.utils import get_cuda_version # from tests.unit_tests.test_utilities import Utils @@ -114,7 +116,7 @@ def setUp(self): self.pg_collection = ProcessGroupCollection.use_mpu_process_groups() self.seed = 46 - self.hidden_size = 2048 + self.hidden_size = 512 self.n_routed_experts = 8 self.acc_steps = 1 @@ -176,6 +178,12 @@ def _build_moe_layer( self.pg_collection, ) + mix_precision_utils.MixPrecisionLayer(moe_layer, dtype="bfloat16") + + for param in moe_layer.parameters(): + if hasattr(param, "main_grad") and param.main_grad is None: + param.main_grad = paddle.zeros_like(param, dtype=paddle.float32) + return moe_layer @staticmethod @@ -241,7 +249,7 @@ def _run_accumulated_forward_backward(self, moe_layer, input_data_list): self._collect_grads(moe_layer), ) - def test_moe_layer_precision(self): + def run_test_moe_layer_precision(self): """Test MoELayer precision: BF16 sonic-moe vs baseline, FP8 vs BF16. Both baseline and sonic layers are built with the same seed. @@ -262,9 +270,14 @@ def test_moe_layer_precision(self): using_sonic_moe=False, moe_deep_gemm=False ) moe_layer_sonic_bf16 = self._build_moe_layer(using_sonic_moe=True) + moe_layer_sonic_bf16.grouped_gemm_experts.sonic_moe_config.enabled = ( + False + ) + moe_layer_sonic_fp8 = self._build_moe_layer( using_sonic_moe=True, fp8="e4m3" ) + moe_layer_sonic_fp8.grouped_gemm_experts.sonic_moe_config.enabled = True input_data_list = [] for step_idx in range(self.acc_steps): @@ -282,6 +295,7 @@ def test_moe_layer_precision(self): moe_layer_sonic_bf16, input_data_list ) ) + moe_layer_sonic_fp8.grouped_gemm_experts.quant_weight() loss_fp8, output_fp8, grads_fp8 = ( self._run_accumulated_forward_backward( moe_layer_sonic_fp8, input_data_list @@ -369,6 +383,54 @@ def test_moe_layer_precision(self): print("All precision checks passed!") + def run_test_z_bf16_recompute_z(self): + sonic_fp8 = self._build_moe_layer(using_sonic_moe=True, fp8="e4m3") + sonic_fp8_recompute_z = self._build_moe_layer( + using_sonic_moe=True, fp8="e4m3" + ) + + sonic_fp8.grouped_gemm_experts.sonic_moe_config.save_z_fp8 = False + sonic_fp8_recompute_z.grouped_gemm_experts.sonic_moe_config.save_z_fp8 = False + sonic_fp8_recompute_z.grouped_gemm_experts.sonic_moe_config.recompute_z = True + # sonic_fp8.grouped_gemm_experts.sonic_moe_config.enabled = True + input_data_list = [] + for step_idx in range(self.acc_steps): + paddle.seed(self.seed + step_idx) + data = paddle.randn( + [2, 64, self.hidden_size], dtype=paddle.bfloat16 + ) + input_data_list.append(data) + + loss, output, grads = self._run_accumulated_forward_backward( + sonic_fp8, input_data_list + ) + loss_recompute, output_recompute, grads_recompute = ( + self._run_accumulated_forward_backward( + sonic_fp8_recompute_z, input_data_list + ) + ) + self.assertEqual(loss, loss_recompute, "Loss not equal.") + self.assertTrue( + np.array_equal(output.numpy(), output_recompute.numpy()), + "Output not equal.", + ) + + common_grads = set(grads) & set(grads_recompute) + self.assertTrue( + common_grads, + "No common FP8 grad tensors found", + ) + for name in sorted(common_grads): + g1 = grads[name] + g2 = grads_recompute[name] + self.assertTrue( + np.array_equal(g1.numpy(), g2.numpy()), "grad not equal." + ) + + def test_all_cases(self): + self.run_test_moe_layer_precision() + self.run_test_z_bf16_recompute_z() + if __name__ == "__main__": unittest.main()