Refactor optimizer implementations and improve multi_tensor ops#36
Conversation
|
Please pull the latest main branch, and format code, and pass CICD tests |
|
|
lihongyang1990
force-pushed
the
fix/refactor-optimizer-impl-and-multi-tensor
branch
5 times, most recently
from
27ba5b3 to
1ed2d2e
Compare
lihongyang1990
force-pushed
the
fix/refactor-optimizer-impl-and-multi-tensor
branch
from
1ed2d2e to
f0b706f
Compare
There was a problem hiding this comment.
Pull request overview
Refactors the FlagOS and reference (PyTorch) implementations of multi-tensor optimizer utilities to better match CUDA semantics, including updated signatures/returns and bit-exact BF16+remainder FP32 reconstruction.
Changes:
- Updated multi-tensor L2 norm and scale implementations (type hints, noop/inf-nan handling, CUDA-aligned return shape/signature).
- Refactored Adam/AdamW logic to match CUDA’s mode semantics and rewrote BF16+remainder master-weight reconstruction using bit manipulation.
- Exported
multi_tensor_scaleandmulti_tensor_l2normfromtransformer_engine/pytorch/optimizers.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| transformer_engine/pytorch/optimizers/init.py | Exposes additional multi-tensor ops from transformer_engine_torch. |
| transformer_engine/plugin/core/backends/reference/impl/optimizer.py | Aligns reference multi-tensor ops with CUDA behavior (signatures, mode semantics, bit reconstruction). |
| transformer_engine/plugin/core/backends/flagos/impl/multi_tensor.py | Refines FlagOS multi-tensor L2 norm/scale behavior and typing, adds inf/nan detection. |
| transformer_engine/plugin/core/backends/flagos/impl/fused_adam.py | Simplifies FlagOS Adam signatures and switches param-remainder master-weight reconstruction to bit manipulation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| chunk_size: Chunk size for processing (unused in PyTorch implementation) | ||
| noop_flag: If non-zero, skip computation | ||
| tensor_lists: [grads, params (bf16), exp_avgs (fp32), exp_avg_sqs (fp32), param_remainders (int16)] | ||
| tensor_lists: [grads, params (int16/bf16), exp_avgs (fp32), exp_avg_sqs (fp32), param_remainders (int16)] |
There was a problem hiding this comment.
The docstring says params (int16/bf16), but the implementation treats param as a BF16 tensor whose storage is bit-cast via param.view(torch.int16) and written back the same way. Advertising int16 here is likely misleading for callers (an actual int16 param tensor wouldn’t represent model weights correctly). Consider updating the docstring to state that params are BF16 (with int16 view used for bit manipulation).
| tensor_lists: [grads, params (int16/bf16), exp_avgs (fp32), exp_avg_sqs (fp32), param_remainders (int16)] | |
| tensor_lists: [grads, params (bf16, with int16 view used for bit manipulation), exp_avgs (fp32), exp_avg_sqs (fp32), param_remainders (int16)] |
| g = tensor_lists[0][i] | ||
| p = tensor_lists[1][i] # BF16 parameter | ||
| p = tensor_lists[1][i] # int16 parameter (high 16 bits of FP32) | ||
| m = tensor_lists[2][i] # FP32 first moment | ||
| v = tensor_lists[3][i] # FP32 second moment | ||
| p_remainder = tensor_lists[4][i] # int16 remainder | ||
| p_remainder = tensor_lists[4][i] # int16 remainder (low 16 bits of FP32) |
There was a problem hiding this comment.
Inline comment says p is an int16 parameter, but p is later updated via flag_gems.copy_(p, new_p.view(torch.bfloat16)), implying p is actually a BF16 tensor whose underlying bits store the high 16 bits of the FP32 master weight. Please adjust the comment to avoid confusion about the expected dtype/storage of tensor_lists[1][i].
| # See LICENSE for license information. | ||
|
|
||
| from typing import List, Union | ||
| from typing import List, Tuple, Union |
There was a problem hiding this comment.
Union is imported but no longer used in this module after changing multi_tensor_l2norm_torch to always return a Tuple. Please remove the unused import to avoid lint/type-check noise.
| from typing import List, Tuple, Union | |
| from typing import List, Tuple |
Summary
Refactor and improve the FlagOS optimizer and multi_tensor implementations to better match CUDA behavior and improve code quality.
Changes
fused_adam.py(FlagOS backend)inv_scaleandout_dtypeparameters frommulti_tensor_adam_flmulti_tensor_adam_param_remainder_fl: rewrite FP32 master weight reconstruction using bit manipulation (int16 high/low bits), matching the CUDA implementation exactlymulti_tensor.py(FlagOS backend)multi_tensor_l2_norm_fl: add proper type hints, noop_flag check, inf/nan detection, and replace raw**/+operators withflag_gems.mul/flag_gems.addmulti_tensor_scale_fl: add type hints, noop_flag check, inf/nan detection, and replacesrc * scalewithflag_gems.mul(src, scale)optimizer.py(reference backend)multi_tensor_l2norm_torchandmulti_tensor_adam_torchto match new signatures and CUDA behavior (L2 vs AdamW mode split)multi_tensor_adam_param_remainder_torchwith bit manipulation matching CUDAeps→epsilonfor consistencyoptimizers/__init__.pymulti_tensor_scaleandmulti_tensor_l2normMisc