From ae56a809345e5ae56a56ba6837c03b7dfbb76ce4 Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Thu, 20 Feb 2025 23:28:03 +0900 Subject: [PATCH 1/2] update typehint Signed-off-by: Masaki Kozuki --- torchao/float8/float8_linear_utils.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/torchao/float8/float8_linear_utils.py b/torchao/float8/float8_linear_utils.py index 3649b741cc..59dea62ce4 100644 --- a/torchao/float8/float8_linear_utils.py +++ b/torchao/float8/float8_linear_utils.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD 3-Clause license found in the # LICENSE file in the root directory of this source tree. import logging -from typing import Callable, Optional +from typing import Callable, List, Optional import torch import torch.distributed as dist @@ -126,7 +126,7 @@ def convert_to_float8_training( module: nn.Module, *, module_filter_fn: Optional[Callable[[nn.Module, str], bool]] = None, - config: Float8LinearConfig = None, + config: Optional[Float8LinearConfig] = None, ) -> nn.Module: """ Swaps `torch.nn.Linear` in `module` with `Float8Linear`. @@ -178,7 +178,10 @@ def get_float8_layers(model: torch.nn.Module): @torch.no_grad() -def sync_float8_amax_and_scale_history(model: torch.nn.Module, fp8_layers=None) -> None: +def sync_float8_amax_and_scale_history( + model: torch.nn.Module, + fp8_layers: Optional[List[Float8Linear]] = None, +) -> None: """ Manages the float8 amax and scale bookkeeping. In detail, it does the following: From b84bc5100e02b6980924f619728015106ae64cf2 Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Mon, 3 Mar 2025 19:58:42 -0800 Subject: [PATCH 2/2] Update float8_linear_utils.py --- torchao/float8/float8_linear_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchao/float8/float8_linear_utils.py b/torchao/float8/float8_linear_utils.py index e38605f0fb..8ea6e2e23a 100644 --- a/torchao/float8/float8_linear_utils.py +++ b/torchao/float8/float8_linear_utils.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD 3-Clause license found in the # LICENSE file in the root directory of this source tree. import logging -from typing import Callable, List, Optional +from typing import Callable, Optional import torch.nn as nn @@ -112,4 +112,4 @@ def convert_to_float8_training( module, from_float, module_filter_fn=module_filter_fn, - ) \ No newline at end of file + )