From 3c833002661c74845b1eeb11e3593ceb6c1d2dd6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=B0=A2=E7=BF=8A=E5=87=A1?= Date: Thu, 20 Aug 2026 21:36:34 +0800 Subject: [PATCH] fix: use torch.amp via compatibility shim to avoid deprecated torch.cuda.amp warnings torch.cuda.amp.autocast has been deprecated since torch 2.4 and torch.cuda.amp.GradScaler since torch 2.3. Both still work but emit a FutureWarning on every use; the replacement device-agnostic APIs live in torch.amp. Add funasr/utils/amp.py, a shim that re-exports torch.amp names when available (torch >= 2.3) and falls back to torch.cuda.amp otherwise. The shim also supplies device_type='cuda' by default for autocast, matching the old torch.cuda.amp.autocast signature, so existing call sites work unchanged. Replace all 37 `from torch.cuda.amp import ...` imports with the shim, and rewrite the 7 `torch.cuda.amp.autocast(...)` call sites to use the imported name so the deprecation path is actually avoided. Co-Authored-By: Claude --- funasr/bin/train.py | 2 +- funasr/bin/train_ds.py | 2 +- funasr/models/bat/model.py | 2 +- funasr/models/bicif_paraformer/model.py | 2 +- funasr/models/campplus/model.py | 2 +- funasr/models/contextual_paraformer/model.py | 2 +- funasr/models/ct_transformer/model.py | 2 +- .../models/ct_transformer_streaming/model.py | 2 +- funasr/models/data2vec/data2vec.py | 2 +- funasr/models/e_paraformer/model.py | 2 +- funasr/models/e_paraformer/pif_predictor.py | 2 +- funasr/models/emotion2vec/model.py | 2 +- funasr/models/fsmn_kws/model.py | 2 +- funasr/models/fsmn_kws_mt/model.py | 2 +- funasr/models/lcbnet/model.py | 2 +- funasr/models/llm_asr/model.py | 14 ++++---- funasr/models/llm_asr_nar/model.py | 2 +- funasr/models/mfcca/e2e_asr_mfcca.py | 2 +- funasr/models/monotonic_aligner/model.py | 2 +- funasr/models/paraformer/cif_predictor.py | 2 +- funasr/models/paraformer/model.py | 2 +- funasr/models/paraformer_streaming/model.py | 2 +- .../models/paraformer_v2_community/model.py | 2 +- funasr/models/sa_asr/e2e_sa_asr.py | 2 +- funasr/models/sanm_kws/model.py | 2 +- funasr/models/sanm_kws_streaming/model.py | 2 +- funasr/models/scama/model.py | 2 +- funasr/models/seaco_paraformer/model.py | 2 +- funasr/models/sense_voice/model.py | 2 +- funasr/models/sond/e2e_diar_sond.py | 2 +- funasr/models/transducer/model.py | 2 +- funasr/models/transformer/model.py | 2 +- funasr/models/uniasr/model.py | 2 +- funasr/models/whisper_lid/model.py | 2 +- funasr/models/xvector/e2e_sv.py | 2 +- funasr/train_utils/trainer.py | 2 +- funasr/train_utils/trainer_ds.py | 4 +-- funasr/utils/amp.py | 36 +++++++++++++++++++ 38 files changed, 80 insertions(+), 44 deletions(-) create mode 100644 funasr/utils/amp.py diff --git a/funasr/bin/train.py b/funasr/bin/train.py index d54cd0bbe..d7c20f8f0 100644 --- a/funasr/bin/train.py +++ b/funasr/bin/train.py @@ -16,7 +16,7 @@ import torch.distributed as dist from omegaconf import DictConfig, OmegaConf -from torch.cuda.amp import autocast, GradScaler +from funasr.utils.amp import autocast, GradScaler from torch.nn.parallel import DistributedDataParallel as DDP from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.algorithms.join import Join diff --git a/funasr/bin/train_ds.py b/funasr/bin/train_ds.py index d1f25973e..1b18fe1b6 100644 --- a/funasr/bin/train_ds.py +++ b/funasr/bin/train_ds.py @@ -15,7 +15,7 @@ import torch.distributed as dist from omegaconf import DictConfig, OmegaConf -from torch.cuda.amp import autocast, GradScaler +from funasr.utils.amp import autocast, GradScaler from torch.nn.parallel import DistributedDataParallel as DDP from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.algorithms.join import Join diff --git a/funasr/models/bat/model.py b/funasr/models/bat/model.py index 587eb56d1..ef88c81e9 100644 --- a/funasr/models/bat/model.py +++ b/funasr/models/bat/model.py @@ -24,7 +24,7 @@ if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/bicif_paraformer/model.py b/funasr/models/bicif_paraformer/model.py index 0a9f86bae..6683aac85 100644 --- a/funasr/models/bicif_paraformer/model.py +++ b/funasr/models/bicif_paraformer/model.py @@ -26,7 +26,7 @@ from funasr.train_utils.device_funcs import to_device if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/campplus/model.py b/funasr/models/campplus/model.py index 12b739137..751c30bcf 100644 --- a/funasr/models/campplus/model.py +++ b/funasr/models/campplus/model.py @@ -26,7 +26,7 @@ if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/contextual_paraformer/model.py b/funasr/models/contextual_paraformer/model.py index 8fdce48d3..4c4df739b 100644 --- a/funasr/models/contextual_paraformer/model.py +++ b/funasr/models/contextual_paraformer/model.py @@ -29,7 +29,7 @@ if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/ct_transformer/model.py b/funasr/models/ct_transformer/model.py index 1ee0bf566..c7beb1ba0 100644 --- a/funasr/models/ct_transformer/model.py +++ b/funasr/models/ct_transformer/model.py @@ -23,7 +23,7 @@ except: pass if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/ct_transformer_streaming/model.py b/funasr/models/ct_transformer_streaming/model.py index c646e2484..06add5f2e 100644 --- a/funasr/models/ct_transformer_streaming/model.py +++ b/funasr/models/ct_transformer_streaming/model.py @@ -16,7 +16,7 @@ if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/data2vec/data2vec.py b/funasr/models/data2vec/data2vec.py index c7386e35d..0ae788c6c 100644 --- a/funasr/models/data2vec/data2vec.py +++ b/funasr/models/data2vec/data2vec.py @@ -22,7 +22,7 @@ from funasr.train_utils.device_funcs import force_gatherable if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/e_paraformer/model.py b/funasr/models/e_paraformer/model.py index d59e1f887..8699032f7 100644 --- a/funasr/models/e_paraformer/model.py +++ b/funasr/models/e_paraformer/model.py @@ -8,7 +8,7 @@ import copy import torch import logging -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from typing import Union, Dict, List, Tuple, Optional from funasr.register import tables diff --git a/funasr/models/e_paraformer/pif_predictor.py b/funasr/models/e_paraformer/pif_predictor.py index 6cd666486..6a7307753 100644 --- a/funasr/models/e_paraformer/pif_predictor.py +++ b/funasr/models/e_paraformer/pif_predictor.py @@ -11,7 +11,7 @@ from funasr.register import tables from funasr.train_utils.device_funcs import to_device from funasr.models.transformer.utils.nets_utils import make_pad_mask -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast @tables.register("predictor_classes", "PifPredictor") diff --git a/funasr/models/emotion2vec/model.py b/funasr/models/emotion2vec/model.py index 3ecace8eb..1166e58d5 100644 --- a/funasr/models/emotion2vec/model.py +++ b/funasr/models/emotion2vec/model.py @@ -23,7 +23,7 @@ logger = logging.getLogger(__name__) if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/fsmn_kws/model.py b/funasr/models/fsmn_kws/model.py index fa6f5b712..03114bd3c 100644 --- a/funasr/models/fsmn_kws/model.py +++ b/funasr/models/fsmn_kws/model.py @@ -6,7 +6,7 @@ import time import torch import logging -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from typing import Union, Dict, List, Tuple, Optional from funasr.register import tables diff --git a/funasr/models/fsmn_kws_mt/model.py b/funasr/models/fsmn_kws_mt/model.py index 3ea002c34..1d7db97e1 100644 --- a/funasr/models/fsmn_kws_mt/model.py +++ b/funasr/models/fsmn_kws_mt/model.py @@ -6,7 +6,7 @@ import time import torch import logging -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from typing import Union, Dict, List, Tuple, Optional from funasr.register import tables diff --git a/funasr/models/lcbnet/model.py b/funasr/models/lcbnet/model.py index cc9aa4e2f..1516090e3 100644 --- a/funasr/models/lcbnet/model.py +++ b/funasr/models/lcbnet/model.py @@ -9,7 +9,7 @@ import time import torch import torch.nn as nn -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from funasr.losses.label_smoothing_loss import LabelSmoothingLoss from funasr.models.ctc.ctc import CTC diff --git a/funasr/models/llm_asr/model.py b/funasr/models/llm_asr/model.py index 798ee103e..5a1ccd5f4 100644 --- a/funasr/models/llm_asr/model.py +++ b/funasr/models/llm_asr/model.py @@ -5,7 +5,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast import re from funasr.models.scama.utils import sequence_mask from funasr.losses.label_smoothing_loss import LabelSmoothingLoss @@ -583,7 +583,7 @@ def forward( batch_size, frames, _ = speech.shape - with torch.cuda.amp.autocast(enabled=False): + with autocast(enabled=False): # audio encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) @@ -618,7 +618,7 @@ def forward( batch_idx, :min_len, : ] - with torch.cuda.amp.autocast( + with autocast( enabled=True if self.llm_dtype != "fp32" else False, dtype=dtype_map[self.llm_dtype] ): labels_ids[labels_ids == -1] = -100 @@ -883,7 +883,7 @@ def inference( llm_dtype = "fp16" if kwargs.get("fp16", False) else llm_dtype llm_dtype = "bf16" if kwargs.get("bf16", False) else llm_dtype - with torch.cuda.amp.autocast( + with autocast( enabled=True if llm_dtype != "fp32" else False, dtype=dtype_map[llm_dtype] ): label = contents["assistant"][0] @@ -1159,7 +1159,7 @@ def forward( batch_size_speech, frames, _ = speech.shape batch_size, token_num = input_ids.shape - with torch.cuda.amp.autocast(enabled=False): + with autocast(enabled=False): # audio encoder encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) @@ -1203,7 +1203,7 @@ def forward( speech_idx += 1 - with torch.cuda.amp.autocast( + with autocast( enabled=True if self.llm_dtype != "fp32" else False, dtype=dtype_map[self.llm_dtype] ): labels_ids[labels_ids == -1] = -100 @@ -1551,7 +1551,7 @@ def inference( llm_dtype = "fp16" if kwargs.get("fp16", False) else llm_dtype llm_dtype = "bf16" if kwargs.get("bf16", False) else llm_dtype - with torch.cuda.amp.autocast( + with autocast( enabled=True if llm_dtype != "fp32" else False, dtype=dtype_map[llm_dtype] ): label = contents["assistant"][-1] diff --git a/funasr/models/llm_asr_nar/model.py b/funasr/models/llm_asr_nar/model.py index b0d4f10d3..d0504f5f6 100644 --- a/funasr/models/llm_asr_nar/model.py +++ b/funasr/models/llm_asr_nar/model.py @@ -5,7 +5,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from funasr.models.scama.utils import sequence_mask from funasr.losses.label_smoothing_loss import LabelSmoothingLoss diff --git a/funasr/models/mfcca/e2e_asr_mfcca.py b/funasr/models/mfcca/e2e_asr_mfcca.py index f6cbb089f..514432bb3 100644 --- a/funasr/models/mfcca/e2e_asr_mfcca.py +++ b/funasr/models/mfcca/e2e_asr_mfcca.py @@ -25,7 +25,7 @@ from funasr.models.base_model import FunASRModel if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/monotonic_aligner/model.py b/funasr/models/monotonic_aligner/model.py index d68e469d1..970ed8d6b 100644 --- a/funasr/models/monotonic_aligner/model.py +++ b/funasr/models/monotonic_aligner/model.py @@ -6,7 +6,7 @@ import time import copy import torch -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from typing import Union, Dict, List, Tuple, Optional from funasr.register import tables diff --git a/funasr/models/paraformer/cif_predictor.py b/funasr/models/paraformer/cif_predictor.py index 93645a8ff..927ae2195 100644 --- a/funasr/models/paraformer/cif_predictor.py +++ b/funasr/models/paraformer/cif_predictor.py @@ -10,7 +10,7 @@ from funasr.register import tables from funasr.train_utils.device_funcs import to_device from funasr.models.transformer.utils.nets_utils import make_pad_mask -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast @tables.register("predictor_classes", "CifPredictor") diff --git a/funasr/models/paraformer/model.py b/funasr/models/paraformer/model.py index 5060d8f20..3f5b01517 100644 --- a/funasr/models/paraformer/model.py +++ b/funasr/models/paraformer/model.py @@ -7,7 +7,7 @@ import copy import torch import logging -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from typing import Union, Dict, List, Tuple, Optional from funasr.register import tables diff --git a/funasr/models/paraformer_streaming/model.py b/funasr/models/paraformer_streaming/model.py index 287ffdfae..69826befa 100644 --- a/funasr/models/paraformer_streaming/model.py +++ b/funasr/models/paraformer_streaming/model.py @@ -26,7 +26,7 @@ if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/paraformer_v2_community/model.py b/funasr/models/paraformer_v2_community/model.py index eb11ede5b..3a0dae5ff 100644 --- a/funasr/models/paraformer_v2_community/model.py +++ b/funasr/models/paraformer_v2_community/model.py @@ -7,7 +7,7 @@ import copy import torch import logging -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from typing import Union, Dict, List, Tuple, Optional from funasr.register import tables diff --git a/funasr/models/sa_asr/e2e_sa_asr.py b/funasr/models/sa_asr/e2e_sa_asr.py index 6d9422d2f..8b0a4fc61 100644 --- a/funasr/models/sa_asr/e2e_sa_asr.py +++ b/funasr/models/sa_asr/e2e_sa_asr.py @@ -29,7 +29,7 @@ from funasr.models.base_model import FunASRModel if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/sanm_kws/model.py b/funasr/models/sanm_kws/model.py index 35a01837f..ecd48332e 100644 --- a/funasr/models/sanm_kws/model.py +++ b/funasr/models/sanm_kws/model.py @@ -6,7 +6,7 @@ import time import torch import logging -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from typing import Union, Dict, List, Tuple, Optional from funasr.register import tables diff --git a/funasr/models/sanm_kws_streaming/model.py b/funasr/models/sanm_kws_streaming/model.py index 3228a0643..43486d97f 100644 --- a/funasr/models/sanm_kws_streaming/model.py +++ b/funasr/models/sanm_kws_streaming/model.py @@ -26,7 +26,7 @@ if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/scama/model.py b/funasr/models/scama/model.py index 76887e5b9..b1c6409a8 100644 --- a/funasr/models/scama/model.py +++ b/funasr/models/scama/model.py @@ -28,7 +28,7 @@ from funasr.models.scama.utils import sequence_mask if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/seaco_paraformer/model.py b/funasr/models/seaco_paraformer/model.py index 720aca087..efa4745a3 100644 --- a/funasr/models/seaco_paraformer/model.py +++ b/funasr/models/seaco_paraformer/model.py @@ -33,7 +33,7 @@ if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/sense_voice/model.py b/funasr/models/sense_voice/model.py index 3ec936c54..c10c663d9 100644 --- a/funasr/models/sense_voice/model.py +++ b/funasr/models/sense_voice/model.py @@ -6,7 +6,7 @@ import torch.nn.functional as F from torch import Tensor from torch import nn -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from funasr.metrics.compute_acc import compute_accuracy, th_accuracy from funasr.losses.label_smoothing_loss import LabelSmoothingLoss from funasr.train_utils.device_funcs import force_gatherable diff --git a/funasr/models/sond/e2e_diar_sond.py b/funasr/models/sond/e2e_diar_sond.py index 44cf6d361..db61dc625 100644 --- a/funasr/models/sond/e2e_diar_sond.py +++ b/funasr/models/sond/e2e_diar_sond.py @@ -29,7 +29,7 @@ from funasr.utils.hinter import hint_once if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/transducer/model.py b/funasr/models/transducer/model.py index 354413161..99fee01f5 100644 --- a/funasr/models/transducer/model.py +++ b/funasr/models/transducer/model.py @@ -23,7 +23,7 @@ if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/models/transformer/model.py b/funasr/models/transformer/model.py index b36864bd8..2cd968773 100644 --- a/funasr/models/transformer/model.py +++ b/funasr/models/transformer/model.py @@ -4,7 +4,7 @@ import time import torch import torch.nn as nn -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from funasr.losses.label_smoothing_loss import LabelSmoothingLoss from funasr.models.ctc.ctc import CTC diff --git a/funasr/models/uniasr/model.py b/funasr/models/uniasr/model.py index b6a767abb..95eea5256 100644 --- a/funasr/models/uniasr/model.py +++ b/funasr/models/uniasr/model.py @@ -6,7 +6,7 @@ import time import torch import logging -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from typing import Union, Dict, List, Tuple, Optional from funasr.register import tables diff --git a/funasr/models/whisper_lid/model.py b/funasr/models/whisper_lid/model.py index e56c01c69..6bd71f1e5 100644 --- a/funasr/models/whisper_lid/model.py +++ b/funasr/models/whisper_lid/model.py @@ -5,7 +5,7 @@ import torch import numpy as np import torch.nn as nn -from torch.cuda.amp import autocast +from funasr.utils.amp import autocast from funasr.losses.label_smoothing_loss import LabelSmoothingLoss from funasr.models.ctc.ctc import CTC diff --git a/funasr/models/xvector/e2e_sv.py b/funasr/models/xvector/e2e_sv.py index c6062b59c..34051903e 100644 --- a/funasr/models/xvector/e2e_sv.py +++ b/funasr/models/xvector/e2e_sv.py @@ -31,7 +31,7 @@ from funasr.models.base_model import FunASRModel if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"): - from torch.cuda.amp import autocast + from funasr.utils.amp import autocast else: # Nothing to do if torch<1.6.0 @contextmanager diff --git a/funasr/train_utils/trainer.py b/funasr/train_utils/trainer.py index 15654cd47..b0493c03f 100644 --- a/funasr/train_utils/trainer.py +++ b/funasr/train_utils/trainer.py @@ -6,7 +6,7 @@ from tqdm import tqdm from datetime import datetime import torch.distributed as dist -from torch.cuda.amp import autocast, GradScaler +from funasr.utils.amp import autocast, GradScaler from contextlib import nullcontext, contextmanager from pathlib import Path diff --git a/funasr/train_utils/trainer_ds.py b/funasr/train_utils/trainer_ds.py index 83cd19a67..9bd7ca31e 100644 --- a/funasr/train_utils/trainer_ds.py +++ b/funasr/train_utils/trainer_ds.py @@ -6,7 +6,7 @@ from tqdm import tqdm from datetime import datetime import torch.distributed as dist -from torch.cuda.amp import autocast, GradScaler +from funasr.utils.amp import autocast, GradScaler from contextlib import nullcontext, contextmanager from pathlib import Path from torch.nn.parallel import DistributedDataParallel as DDP @@ -32,7 +32,7 @@ def maybe_autocast(dtype=None, use_deepspeed=False): use_deepspeed: TODO. """ if use_deepspeed: - with torch.cuda.amp.autocast(enabled=True, dtype=dtype, cache_enabled=False): + with autocast(enabled=True, dtype=dtype, cache_enabled=False): yield else: if dtype == torch.float16 or dtype == torch.bfloat16: diff --git a/funasr/utils/amp.py b/funasr/utils/amp.py new file mode 100644 index 000000000..5656a1d18 --- /dev/null +++ b/funasr/utils/amp.py @@ -0,0 +1,36 @@ +"""Compatibility shim for torch.amp (autocast / GradScaler). + +``torch.cuda.amp.autocast`` has been deprecated since torch 2.4 and +``torch.cuda.amp.GradScaler`` since torch 2.3 — both still work but emit +a ``FutureWarning`` on every use, and the replacement device-agnostic +APIs live in ``torch.amp`` (``autocast('cuda', ...)`` since 2.0, +``GradScaler('cuda', ...)`` since 2.3). + +This module re-exports the non-deprecated ``torch.amp`` names when the +installed torch provides them, and falls back to ``torch.cuda.amp`` on +older torch versions. Import it instead of ``torch.cuda.amp`` directly: + + from funasr.utils.amp import autocast, GradScaler + +Note: ``torch.amp.autocast`` requires a ``device_type`` argument (e.g. +``'cuda'``) that the deprecated ``torch.cuda.amp.autocast`` did not, so +the shim supplies ``device_type="cuda"`` by default for callers that use +the old signature (``with autocast(enabled=..., dtype=...)``). +""" + +import torch + +if hasattr(torch.amp, "autocast") and hasattr(torch.amp, "GradScaler"): + from torch.amp import autocast as _amp_autocast + from torch.amp import GradScaler + + def autocast(*args, **kwargs): + """torch.amp.autocast with device_type defaulting to "cuda".""" + if not args and "device_type" not in kwargs: + kwargs["device_type"] = "cuda" + return _amp_autocast(*args, **kwargs) + +else: + from torch.cuda.amp import autocast, GradScaler + +__all__ = ["autocast", "GradScaler"]