github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[fp8] add fallback and make compile option configurable (#6092)

pull/6094/head
Hongxin Liu 2024-10-18 13:55:31 +08:00 committed by GitHub
parent 3b1d7d1ae8
commit 5ddad486ca
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/quantization/fp8.py
View File

@ -8,6 +8,8 @@ import torch.nn.functional as F
from packaging.version import Version
from torch.distributed import ReduceOp
from .fp8_config import dynamic_kernel
SUPPORT_TORCH_COMPILE = Version(torch.__version__) >= Version("2.4.0")
SCALE_BYTES = 4
try:
@ -832,11 +834,13 @@ class _LinearFp8(torch.autograd.Function):
return x_grad.reshape(ctx.x_shape), w_grad, bias_grad
@torch.compile(mode="max-autotune-no-cudagraphs", disable=not SUPPORT_TORCH_COMPILE, dynamic=False)
@torch.compile(mode="max-autotune-no-cudagraphs", disable=not SUPPORT_TORCH_COMPILE, dynamic=dynamic_kernel)
def _linear_fp8(input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor:
return _LinearFp8.apply(input, weight, bias)
def linear_fp8(input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor:
if input.shape[-1] % 16 != 0 or np.prod(input.shape[:-1]) % 16 != 0:
return F.linear(input, weight, bias)
out = _linear_fp8(input, weight, bias)
return out

1
colossalai/quantization/fp8_config.py Normal file
View File

@ -0,0 +1 @@
dynamic_kernel: bool = False