已合并
perf: defer mxfp8 backward quantization #3852
guihaowen666创建于 8月6日
perf: defer mxfp8 backward quantization #3852
已合并
共 3 个文件变更+79-28
| @@ -189,15 +189,24 @@ class MXFP8GMMFunction(BaseGMMFunction): | |||
| 189 | def op_forward(cls, ctx, x, weight, group_list, group_list_type=0, bias=None, reuse_identity=None): | 189 | def op_forward(cls, ctx, x, weight, group_list, group_list_type=0, bias=None, reuse_identity=None): |
| 190 | qdtype = get_quant_dtype() | 190 | qdtype = get_quant_dtype() |
| 191 | x_mxfp8, x_scale = torch_npu.npu_dynamic_mx_quant(x, axis=-1, dst_type=qdtype.x) | 191 | x_mxfp8, x_scale = torch_npu.npu_dynamic_mx_quant(x, axis=-1, dst_type=qdtype.x) |
| 192 | - weight_col_mxfp8, weight_col_scale, weight_row_mxfp8, weight_row_scale = reuse_or_quantize( | 192 | + if getattr(get_args(), 'mxfp8_defer_backward_quant', False): |
| 193 | - weight, | 193 | + weight_row_mxfp8, weight_row_scale = reuse_or_quantize( |
| 194 | - TensorKey.weight, | 194 | + weight, |
| 195 | - torch_npu.npu_dynamic_mx_quant_with_dual_axis, | 195 | + TensorKey.weight, |
| 196 | - op_name="npu_dynamic_mx_quant_with_dual_axis", | 196 | + torch_npu.npu_dynamic_mx_quant, |
| 197 | - reuse_identity=reuse_identity, | 197 | + axis=-2, |
| 198 | - dst_type=qdtype.w, | 198 | + dst_type=qdtype.w, |
| 199 | - ) | 199 | + ) |
| 200 | - ctx.w_quant = (weight_col_mxfp8, weight_col_scale) | 200 | + else: |
 | |||
| 201 | + weight_col_mxfp8, weight_col_scale, weight_row_mxfp8, weight_row_scale = reuse_or_quantize( | ||
| 202 | + weight, | ||
| 203 | + TensorKey.weight, | ||
| 204 | + torch_npu.npu_dynamic_mx_quant_with_dual_axis, | ||
| 205 | + op_name="npu_dynamic_mx_quant_with_dual_axis", | ||
| 206 | + reuse_identity=reuse_identity, | ||
| 207 | + dst_type=qdtype.w, | ||
| 208 | + ) | ||
| 209 | + ctx.w_quant = (weight_col_mxfp8, weight_col_scale) | ||
| 201 | return torch_npu.npu_grouped_matmul( | 210 | return torch_npu.npu_grouped_matmul( |
| 202 | [x_mxfp8], | 211 | [x_mxfp8], |
| 203 | [weight_row_mxfp8], | 212 | [weight_row_mxfp8], |
| @@ -217,7 +226,18 @@ class MXFP8GMMFunction(BaseGMMFunction): | |||
| 217 | def op_dx(cls, ctx, grad, weight, group_list, group_list_type=0, bias=None): | 226 | def op_dx(cls, ctx, grad, weight, group_list, group_list_type=0, bias=None): |
| 218 | qdtype = get_quant_dtype() | 227 | qdtype = get_quant_dtype() |
| 219 | grad_mxfp8, grad_scale = torch_npu.npu_dynamic_mx_quant(grad, axis=-1, dst_type=qdtype.grads) | 228 | grad_mxfp8, grad_scale = torch_npu.npu_dynamic_mx_quant(grad, axis=-1, dst_type=qdtype.grads) |
| 220 | - weight_mxfp8, weight_scale = ctx.w_quant | 229 | + |
| 230 | + if hasattr(ctx, 'w_quant'): | ||
| 231 | + weight_mxfp8, weight_scale = ctx.w_quant | ||
| 232 | + else: | ||
| 233 | + weight_mxfp8, weight_scale = reuse_or_quantize( | ||
| 234 | + weight, | ||
| 235 | + TensorKey.weight, | ||
| 236 | + torch_npu.npu_dynamic_mx_quant, | ||
| 237 | + axis=-1, | ||
| 238 | + dst_type=qdtype.w, | ||
| 239 | + ) | ||
| 240 | + | ||
| 221 | return torch_npu.npu_grouped_matmul( | 241 | return torch_npu.npu_grouped_matmul( |
| 222 | [grad_mxfp8], | 242 | [grad_mxfp8], |
| 223 | [rearrange(weight_mxfp8, 'n h f -> n f h')], | 243 | [rearrange(weight_mxfp8, 'n h f -> n f h')], |
| @@ -82,6 +82,13 @@ class TransformerEngineBasicFeature(MindSpeedFeature): | |||
| 82 | help="Reuse quantized FP8 weight tensors within one optimizer step.", | 82 | help="Reuse quantized FP8 weight tensors within one optimizer step.", |
| 83 | ) | 83 | ) |
| 84 | 84 | ||
| 85 | + group.add_argument( | ||
| 86 | + "--mxfp8-defer-backward-quant", | ||
| 87 | + action="store_true", | ||
| 88 | + default=False, | ||
| 89 | + help="Defer backward-requierd MXFP8 quantization until backward to reduce saved-tensor memory at the cost of additional backward compute.", | ||
| 90 | + ) | ||
| 91 | + | ||
| 85 | def validate_args(self, args): | 92 | def validate_args(self, args): |
| 86 | if args.fp8 and args.transformer_impl == 'local': | 93 | if args.fp8 and args.transformer_impl == 'local': |
| 87 | raise AssertionError('FP8 just support TE implement.') | 94 | raise AssertionError('FP8 just support TE implement.') |
| @@ -105,6 +112,9 @@ class TransformerEngineBasicFeature(MindSpeedFeature): | |||
| 105 | if getattr(args, "fp8_reuse_quantized_weight", False) and not args.fp8: | 112 | if getattr(args, "fp8_reuse_quantized_weight", False) and not args.fp8: |
| 106 | raise ValueError("fp8_reuse_quantized_weight is only valid when FP8 training is enabled") | 113 | raise ValueError("fp8_reuse_quantized_weight is only valid when FP8 training is enabled") |
| 107 | 114 | ||
| 115 | + if getattr(args, "mxfp8_defer_backward_quant", False) and not args.fp8: | ||
| 116 | + raise ValueError("mxfp8_defer_backward_quant is only valid when FP8 training is enabled") | ||
| 117 | + | ||
| 108 | def pre_register_patches(self, patch_manager, args): | 118 | def pre_register_patches(self, patch_manager, args): |
| 109 | patch_manager.register_patch( | 119 | patch_manager.register_patch( |
| 110 | 'transformer_engine.pytorch.tensor.QuantizedTensor', torch.nn.Module, create_dummy=True | 120 | 'transformer_engine.pytorch.tensor.QuantizedTensor', torch.nn.Module, create_dummy=True |
| @@ -8,6 +8,7 @@ from mindspeed.te.pytorch.fp8.constants import TensorKey | |||
| 8 | from mindspeed.te.pytorch.fp8.recipes.recipe import Recipe, RecipeScaling | 8 | from mindspeed.te.pytorch.fp8.recipes.recipe import Recipe, RecipeScaling |
| 9 | from mindspeed.te.pytorch.fp8.reuse import reuse_or_quantize | 9 | from mindspeed.te.pytorch.fp8.reuse import reuse_or_quantize |
| 10 | from mindspeed.te.pytorch.utils import view_as_n_dim, get_quant_dtype | 10 | from mindspeed.te.pytorch.utils import view_as_n_dim, get_quant_dtype |
| 11 | +from mindspeed.args_utils import get_full_args as get_args | ||
| 11 | 12 | ||
| 12 | 13 | ||
| 13 | class MXFP8ScalingRecipe(Recipe): | 14 | class MXFP8ScalingRecipe(Recipe): |
| @@ -66,15 +67,15 @@ class MXFP8BlockScaling(RecipeScaling): | |||
| 66 | 67 | ||
| 67 | 68 | ||
| 68 | class MXFP8MatMul(torch.autograd.Function): | 69 | class MXFP8MatMul(torch.autograd.Function): |
| 69 | - | ||
| 70 | 70 | ||
| 71 | def forward(ctx, x: torch.Tensor, weight: torch.Tensor, need_grad: bool = True): | 71 | def forward(ctx, x: torch.Tensor, weight: torch.Tensor, need_grad: bool = True): |
| 72 | qdtype = get_quant_dtype() | 72 | qdtype = get_quant_dtype() |
| 73 | x_2d = view_as_n_dim(x) | 73 | x_2d = view_as_n_dim(x) |
| 74 | ctx.output_dtype = x.dtype | 74 | ctx.output_dtype = x.dtype |
| 75 | - if need_grad: | 75 | + if need_grad and not getattr(get_args(), 'mxfp8_defer_backward_quant', False): |
| 76 | - x_quant, x_scale, ctx.x, ctx.x_scale = \ | 76 | + x_quant, x_scale, ctx.x, ctx.x_scale = torch_npu.npu_dynamic_mx_quant_with_dual_axis( |
| 77 | - torch_npu.npu_dynamic_mx_quant_with_dual_axis(x_2d, dst_type=qdtype.x) | 77 | + x_2d, dst_type=qdtype.x |
| 78 | + ) | ||
| 78 | w_quant, w_scale, ctx.w, ctx.w_scale = reuse_or_quantize( | 79 | w_quant, w_scale, ctx.w, ctx.w_scale = reuse_or_quantize( |
| 79 | weight, | 80 | weight, |
| 80 | TensorKey.weight, | 81 | TensorKey.weight, |
| @@ -91,10 +92,16 @@ class MXFP8MatMul(torch.autograd.Function): | |||
| 91 | dst_type=qdtype.w, | 92 | dst_type=qdtype.w, |
| 92 | ) | 93 | ) |
| 93 | ctx.save_for_backward(x, weight) | 94 | ctx.save_for_backward(x, weight) |
| 94 | - output = torch_npu.npu_quant_matmul(x_quant, w_quant.t(), w_scale.transpose(0, 1), | 95 | + output = torch_npu.npu_quant_matmul( |
| 95 | - pertoken_scale=x_scale, | 96 | + x_quant, |
| 96 | - output_dtype=x.dtype, scale_dtype=torch_npu.float8_e8m0fnu, | 97 | + w_quant.t(), |
| 97 | - pertoken_scale_dtype=torch_npu.float8_e8m0fnu, group_sizes=[1, 1, 32]) | 98 | + w_scale.transpose(0, 1), |
| 99 | + pertoken_scale=x_scale, | ||
| 100 | + output_dtype=x.dtype, | ||
| 101 | + scale_dtype=torch_npu.float8_e8m0fnu, | ||
| 102 | + pertoken_scale_dtype=torch_npu.float8_e8m0fnu, | ||
| 103 | + group_sizes=[1, 1, 32], | ||
| 104 | + ) | ||
| 98 | if len(x.shape) != 2: | 105 | if len(x.shape) != 2: |
| 99 | output = output.reshape(*x.shape[:-1], *output.shape[1:]) | 106 | output = output.reshape(*x.shape[:-1], *output.shape[1:]) |
| 100 | if weight.requires_grad: | 107 | if weight.requires_grad: |
| @@ -104,8 +111,9 @@ class MXFP8MatMul(torch.autograd.Function): | |||
| 104 | 111 | ||
| 105 | def backward(ctx, grads: torch.Tensor): | 112 | def backward(ctx, grads: torch.Tensor): |
| 106 | qdtype = get_quant_dtype() | 113 | qdtype = get_quant_dtype() |
| 107 | - grads_dx, grads_dx_scale, grads_dw, grads_dw_scale = \ | 114 | + grads_dx, grads_dx_scale, grads_dw, grads_dw_scale = torch_npu.npu_dynamic_mx_quant_with_dual_axis( |
| 108 | - torch_npu.npu_dynamic_mx_quant_with_dual_axis(view_as_n_dim(grads), dst_type=qdtype.grads) | 115 | + view_as_n_dim(grads), dst_type=qdtype.grads |
| 116 | + ) | ||
| 109 | 117 | ||
| 110 | if hasattr(ctx, 'x'): | 118 | if hasattr(ctx, 'x'): |
| 111 | x_quant, x_scale, w_quant, w_scale = ctx.x, ctx.x_scale, ctx.w, ctx.w_scale | 119 | x_quant, x_scale, w_quant, w_scale = ctx.x, ctx.x_scale, ctx.w, ctx.w_scale |
| @@ -120,14 +128,27 @@ class MXFP8MatMul(torch.autograd.Function): | |||
| 120 | ) | 128 | ) |
| 121 | x_quant, x_scale = torch_npu.npu_dynamic_mx_quant(view_as_n_dim(x), axis=-2, dst_type=qdtype.x) | 129 | x_quant, x_scale = torch_npu.npu_dynamic_mx_quant(view_as_n_dim(x), axis=-2, dst_type=qdtype.x) |
| 122 | 130 | ||
| 123 | - dx = torch_npu.npu_quant_matmul(grads_dx, w_quant, w_scale, | 131 | + dx = torch_npu.npu_quant_matmul( |
| 124 | - pertoken_scale=grads_dx_scale, | 132 | + grads_dx, |
| 125 | - output_dtype=ctx.output_dtype, scale_dtype=torch_npu.float8_e8m0fnu, | 133 | + w_quant, |
| 126 | - pertoken_scale_dtype=torch_npu.float8_e8m0fnu, group_sizes=[1, 1, 32]) | 134 | + w_scale, |
| 135 | + pertoken_scale=grads_dx_scale, | ||
| 136 | + output_dtype=ctx.output_dtype, | ||
| 137 | + scale_dtype=torch_npu.float8_e8m0fnu, | ||
| 138 | + pertoken_scale_dtype=torch_npu.float8_e8m0fnu, | ||
| 139 | + group_sizes=[1, 1, 32], | ||
| 140 | + ) | ||
| 127 | if len(grads.shape) != 2: | 141 | if len(grads.shape) != 2: |
| 128 | dx = dx.reshape(*grads.shape[:-1], *dx.shape[1:]) | 142 | dx = dx.reshape(*grads.shape[:-1], *dx.shape[1:]) |
| 129 | 143 | ||
| 130 | - dw = torch_npu.npu_quant_matmul(grads_dw.t(), x_quant, x_scale, pertoken_scale=grads_dw_scale.transpose(0, 1), | 144 | + dw = torch_npu.npu_quant_matmul( |
| 131 | - output_dtype=ctx.output_dtype, scale_dtype=torch_npu.float8_e8m0fnu, | 145 | + grads_dw.t(), |
| 132 | - pertoken_scale_dtype=torch_npu.float8_e8m0fnu, group_sizes=[1, 1, 32]) | 146 | + x_quant, |
| 133 | - return dx, dw, None, None, None | 147 | + x_scale, |
| 148 | + pertoken_scale=grads_dw_scale.transpose(0, 1), | ||
| 149 | + output_dtype=ctx.output_dtype, | ||
| 150 | + scale_dtype=torch_npu.float8_e8m0fnu, | ||
| 151 | + pertoken_scale_dtype=torch_npu.float8_e8m0fnu, | ||
| 152 | + group_sizes=[1, 1, 32], | ||
| 153 | + ) | ||
| 154 | + return dx, dw, None, None, None | ||
🔴 Critical
变更行：第 192–211 行。
当
fp8_quant_fb_split为 True 时（第 192 行分支），只定义了weight_row_mxfp8和weight_row_scale（单轴量化，axis=-2），未定义weight_col_mxfp8和weight_col_scale。但第 211 行ctx.w_quant = (weight_col_mxfp8, weight_col_scale)无条件执行，且在第 192 行分支中这些变量从未赋值，会直接抛出NameError，导致所有启用该特性的训练立即崩溃。同时，
fp8_quant_fb_split的设计意图是正向不保存量化结果、反向重新量化（op_dx中通过hasattr(ctx, 'w_quant')判断）；因此第 211 行应只在else分支（双轴量化）中执行，否则会同时破坏正确性与设计意图。建议：将第 211 行移入 else 分支（第 210 行之后），或在该行外加
if not getattr(get_args(), 'fp8_quant_fb_split', False):守卫。最简单的方式是删除第 211 行，并将第 210 行的ctx.w_quant = (weight_col_mxfp8, weight_col_scale)作为 else 分支的唯一赋值点。