[WIP] Add Support for msamp FP4 Quantization and Unit Test #203
Conversation
msamp/operators/fp4_quant/setup.py
Outdated
| ) | ||
| ], | ||
| cmdclass={'build_ext': cpp_extension.BuildExtension} | ||
| ) |
There was a problem hiding this comment.
Please add newline at the end of the file
There was a problem hiding this comment.
The same for other files
msamp/operators/fp4_quant/quant.cu
Outdated
| __nv_bfloat16* output_data = reinterpret_cast<__nv_bfloat16*>(output.data_ptr<at::BFloat16>()); | ||
|
|
||
| const int threadsPerBlock = HIP_GET_NUM_THREADS(size); // 512 | ||
| // const int blocks = HIP_GET_BLOCKS(size, threadsPerBlock); // max grid num: HIP_MAX_GRID_NUM = 65535 |
There was a problem hiding this comment.
Please remove all unused code like this
msamp/common/dtype/floating.py
Outdated
| @@ -15,24 +15,29 @@ class Floating: | |||
| qfp_max: dict = {} | |||
|
|
|||
| @staticmethod | |||
| def _get_fp_max(exp, man, inf_existed=True): | |||
| def _get_fp_max(exp, man, inf_existed=True, nan_existed=True): | |||
There was a problem hiding this comment.
Seems we don't need to revise this file, because parameter nan_existed is never set to False in new code
msamp/megatron/layers.py
Outdated
| @@ -13,6 +13,14 @@ | |||
| from msamp.common.tensor import ScalingTensor | |||
| from msamp.operators.gemm import Gemm | |||
|
|
|||
| import os | |||
|
|
|||
| USE_W_SIMU_FP4 = bool(int(os.getenv('USE_W_SIMU_FP4', 0))) | |||
There was a problem hiding this comment.
Shall we make the naming clearer? Like using WEIGHT and ACTIVATION instead of W and A, and use SIMULATION/SIMULATE instead of SIMU. Also we need to add MSAMP_ prefix otherwise it's not bounded in MS-AMP project.
Makefile
Outdated
| @@ -24,4 +24,5 @@ lint: cpplint mdlint | |||
| postinstall: | |||
| cd msamp/operators/dist_op && bash build.sh && cd - | |||
| cd msamp/operators/arithmetic && pip install -v . && cd - | |||
| cd msamp/operators/fp4_quant && pip install -v . && cd - | |||
There was a problem hiding this comment.
Please just use quantize or quantization for naming.
msamp/megatron/layers.py
Outdated
| @@ -175,6 +196,9 @@ def backward(ctx, grad_output): | |||
| wgrad_qtype, | |||
| use_split_accumulator=True, | |||
| ) | |||
| if USE_W_DIFFERENTIABLE_GRADIENT_ESTIMATOR: | |||
| scaled_w = ctx.saved_tensors[0] | |||
| grad_weight.mul_(FP4_QUANT.apply_DGE_item(scaled_w, k=5.0, power_clamp_max=3.0)) | |||
There was a problem hiding this comment.
Let's just eliminate explicit argument assignments for k and power_clamp_max here, if their default values are already provided.
| class FP4_QUANT: | ||
| """FP4 Quantization operator.""" | ||
| @staticmethod | ||
| def apply_DGE_item( |
There was a problem hiding this comment.
What does DGE mean here?
There was a problem hiding this comment.
Please add full term in function comment
|
|
||
|
|
||
| @staticmethod | ||
| def quantize_simu_fp4_in_bf16( |
There was a problem hiding this comment.
simulate or simulation instead of simu?
There was a problem hiding this comment.
Please add MIT license head
Work in Progress, Please donot Merge
This PR is for new features regarding to FP4 quantization for MS-AMP library.
Working items: