mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-13 05:01:44 +00:00
ae02d4e4f7
* [bf16] add bf16 support for fused adam (#3844) * [bf16] fused adam kernel support bf16 * [test] update fused adam kernel test * [test] update fused adam test * [bf16] cpu adam and hybrid adam optimizers support bf16 (#3860) * [bf16] implement mixed precision mixin and add bf16 support for low level zero (#3869) * [bf16] add mixed precision mixin * [bf16] low level zero optim support bf16 * [text] update low level zero test * [text] fix low level zero grad acc test * [bf16] add bf16 support for gemini (#3872) * [bf16] gemini support bf16 * [test] update gemini bf16 test * [doc] update gemini docstring * [bf16] add bf16 support for plugins (#3877) * [bf16] add bf16 support for legacy zero (#3879) * [zero] init context support bf16 * [zero] legacy zero support bf16 * [test] add zero bf16 test * [doc] add bf16 related docstring for legacy zero
24 lines
504 B
Python
24 lines
504 B
Python
import torch
|
|
from torch import Tensor
|
|
|
|
from .base import MixedPrecisionMixin
|
|
|
|
|
|
class BF16MixedPrecisionMixin(MixedPrecisionMixin):
|
|
dtype = torch.bfloat16
|
|
|
|
def pre_backward(self, loss: Tensor) -> Tensor:
|
|
return loss
|
|
|
|
def pre_backward_by_grad(self, tensor: Tensor, grad: Tensor) -> Tensor:
|
|
return grad
|
|
|
|
def should_skip_step(self) -> bool:
|
|
return False
|
|
|
|
def pre_zero_grad(self) -> None:
|
|
pass
|
|
|
|
def get_grad_div_scale(self) -> float:
|
|
return 1.0
|