Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Aug 27, 2024
1 parent 5f2b65b commit 5c649e9
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 3 deletions.
6 changes: 3 additions & 3 deletions transformer_engine/pytorch/fp8.py
Original file line number Diff line number Diff line change
Expand Up @@ -485,9 +485,9 @@ def restore_fp8_meta_tensors(fp8_meta: Dict[str, Any]) -> None:

@contextmanager
def fp8_model_init(
enabled: bool = True,
preserve_high_precision_init_val: bool = False,
) -> None:
enabled: bool = True,
preserve_high_precision_init_val: bool = False,
) -> None:
"""
Context manager for FP8 initialization of parameters.
Expand Down
1 change: 1 addition & 0 deletions transformer_engine/pytorch/module/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -882,6 +882,7 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None:
# a parameter so we always re-apply it just for extra safety.
param = torch.nn.Parameter(param)
if self.primary_weights_in_fp8 and self.preserve_high_precision_init_val:

def get(self):
if hasattr(self, "_high_precision_init_val"):
return self._high_precision_init_val
Expand Down

0 comments on commit 5c649e9

Please sign in to comment.