fix(contrib): call torch.cuda.is_available() in FP16_Optimizer guard

apex/__init__.py

@@ -11,11 +11,9 @@
 # so they expect those backends to be available, but for some reason they actually aren't
 # available (for example because they built improperly in a way that isn't revealed until
 # load time) the error message is timely and visible.
-from . import optimizers
-from . import normalization
+from . import normalization, optimizers
 
-
-__all__ = ["optimizers", "normalization"]
+__all__ = ["normalization", "optimizers"]
 
 
 def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int) -> bool:

apex/_autocast_utils.py

@@ -1,8 +1,8 @@
-from typing import Optional, Sequence
+from collections.abc import Sequence
+from typing import Optional
 
 import torch
 
-
 __all__ = ["_cast_if_autocast_enabled"]
 
 
@@ -12,7 +12,7 @@ def _get_autocast_dtypes() -> Sequence[torch.dtype]:
     return [torch.half]
 
 
-def _get_current_dtype(dtype: Optional[torch.dtype] = None) -> torch.dtype:
+def _get_current_dtype(dtype: torch.dtype | None = None) -> torch.dtype:
     if not torch.is_autocast_enabled():
         return torch.float or dtype
     else:

apex/contrib/bottleneck/__init__.py

@@ -1,7 +1,7 @@
 from .bottleneck import Bottleneck, SpatialBottleneck
 from .halo_exchangers import (
-    HaloExchangerNoComm,
     HaloExchangerAllGather,
-    HaloExchangerSendRecv,
+    HaloExchangerNoComm,
     HaloExchangerPeer,
+    HaloExchangerSendRecv,
 )

apex/contrib/bottleneck/bottleneck.py

@@ -1,12 +1,11 @@
 import functools as func
 
+import fast_bottleneck
+import nccl_p2p_cuda as inc
 import torch
 from torch import nn
 
 from apex import check_cudnn_version_and_warn
-import fast_bottleneck
-import nccl_p2p_cuda as inc
-
 
 assert check_cudnn_version_and_warn(__name__, 8400)
 
@@ -35,7 +34,7 @@ class FrozenBatchNorm2d(torch.jit.ScriptModule):
     """
 
     def __init__(self, n):
-        super(FrozenBatchNorm2d, self).__init__()
+        super().__init__()
         self.register_buffer("weight", torch.ones(n))
         self.register_buffer("bias", torch.zeros(n))
         self.register_buffer("running_mean", torch.zeros(n))
@@ -171,7 +170,7 @@ def __init__(
         use_cudnn=False,
         explicit_nhwc=False,
     ):
-        super(Bottleneck, self).__init__()
+        super().__init__()
         if groups != 1:
             raise RuntimeError("Only support groups == 1")
         if dilation != 1:
@@ -224,8 +223,6 @@ def __init__(
                 with torch.no_grad():
                     p.data = p.data.permute(0, 2, 3, 1).contiguous()
 
-        return
-
     # Returns single callable that recomputes scale and bias for all frozen batch-norms.
     # This method must be called before cuda graphing.
     # The callable it returns can be called anytime.
@@ -851,7 +848,7 @@ def __init__(
         explicit_nhwc=False,
         spatial_parallel_args=None,
     ):
-        super(SpatialBottleneck, self).__init__()
+        super().__init__()
         if groups != 1:
             raise RuntimeError("Only support groups == 1")
         if dilation != 1:
@@ -911,7 +908,6 @@ def __init__(
             self.spatial_parallel_args = (1, 0, None, None, 0, False)
         else:
             self.spatial_parallel_args = spatial_parallel_args
-        return
 
     # Returns single callable that recomputes scale and bias for all frozen batch-norms.
     # This method must be called before cuda graphing.

apex/contrib/bottleneck/halo_exchangers.py

@@ -1,13 +1,13 @@
-import torch
 import nccl_p2p_cuda as inc
 import peer_memory_cuda as pm
+import torch
 
 
 # Communication free halo exchanger.
 # NB! This halo exchanger does not exchange halos with neighbors as it should, it merely swaps the inputs
 # NB! This is only useful for performance testing.
 # NB! Do not use for actual production runs
-class HaloExchanger(object):
+class HaloExchanger:
     def __init__(self, ranks, rank_in_group):
         self.stream1 = torch.cuda.Stream()
         self.stream2 = torch.cuda.Stream()
@@ -27,7 +27,7 @@ def __init__(self, ranks, rank_in_group):
 
 class HaloExchangerNoComm(HaloExchanger):
     def __init__(self, ranks, rank_in_group):
-        super(HaloExchangerNoComm, self).__init__(ranks, rank_in_group)
+        super().__init__(ranks, rank_in_group)
 
     def left_right_halo_exchange(
         self,
@@ -45,7 +45,7 @@ def left_right_halo_exchange(
 
 class HaloExchangerAllGather(HaloExchanger):
     def __init__(self, ranks, rank_in_group, comm):
-        super(HaloExchangerAllGather, self).__init__(ranks, rank_in_group)
+        super().__init__(ranks, rank_in_group)
         # self.comm must be NCCL process_group created with torch.distributed.new_group(ranks=ranks)
         self.comm = comm
 
@@ -94,7 +94,7 @@ def left_right_halo_exchange(
 
 class HaloExchangerSendRecv(HaloExchanger):
     def __init__(self, ranks, rank_in_group):
-        super(HaloExchangerSendRecv, self).__init__(ranks, rank_in_group)
+        super().__init__(ranks, rank_in_group)
         nccl_id = inc.get_unique_nccl_id(1).cuda()
         torch.distributed.broadcast(nccl_id, 0)
         nccl_id = nccl_id.cpu()
@@ -145,7 +145,7 @@ def left_right_halo_exchange(
 
 class HaloExchangerPeer(HaloExchanger):
     def __init__(self, ranks, rank_in_group, peer_pool, explicit_nhwc, numSM=0):
-        super(HaloExchangerPeer, self).__init__(ranks, rank_in_group)
+        super().__init__(ranks, rank_in_group)
         self.diagnostics = False
         self.explicit_nhwc = explicit_nhwc
         self.numSM = numSM

apex/contrib/bottleneck/test.py

@@ -29,7 +29,7 @@
         print("[DEBUG] ref dx :", d_grad.sum().item())
         # print wgrad. we don't need to reset since later cpp print before accumulation
         for i, w in enumerate(model.w_conv):
-            print("[DEBUG] ref wgrad{} :".format(i + 1), w.grad.sum().item())
+            print(f"[DEBUG] ref wgrad{i + 1} :", w.grad.sum().item())
 
     wgrads = []
     for w in model.w_conv:
@@ -56,7 +56,7 @@
     )
     for i, (w, wgrad) in enumerate(zip(model.w_conv, wgrads)):
         print(
-            "max error wgrad{}:".format(i + 1),
+            f"max error wgrad{i + 1}:",
             (wgrad - w.grad.float()).abs().max().item(),
             "max elem:",
             wgrad.abs().max().item(),
@@ -104,7 +104,7 @@
     )
     for i, (w, wgrad) in enumerate(zip(nhwc_model.w_conv, wgrads)):
         print(
-            "max error wgrad{}:".format(i + 1),
+            f"max error wgrad{i + 1}:",
             (wgrad - w.grad.float()).abs().max().item(),
             "max elem:",
             wgrad.abs().max().item(),

apex/contrib/clip_grad/clip_grad.py

@@ -1,10 +1,12 @@
-from typing import Union, Iterable
+from collections.abc import Iterable
+from typing import Union
 
 import torch
 
 _kernel_import_succeeded = False
 try:
     import amp_C
+
     from apex.multi_tensor_apply import multi_tensor_applier
 
     _kernel_import_succeeded = True

apex/contrib/conv_bias_relu/__init__.py

@@ -1,6 +1,6 @@
 from .conv_bias_relu import (
-    ConvBiasReLU,
     ConvBias,
     ConvBiasMaskReLU,
+    ConvBiasReLU,
     ConvFrozenScaleBiasReLU,
 )

apex/contrib/conv_bias_relu/conv_bias_relu.py

@@ -1,7 +1,7 @@
+import fused_conv_bias_relu
 import torch
 
 from apex import check_cudnn_version_and_warn
-import fused_conv_bias_relu
 
 check_cudnn_version_and_warn(__name__, 8400)
 

apex/contrib/csrc/group_norm_v2/generate_gn_cuda_inst.py

@@ -1,6 +1,5 @@
 import pathlib
 
-
 hw_c_list = [
     (8 * 8, 1280),
     (8 * 8, 2560),

apex/contrib/cudnn_gbn/batch_norm.py

@@ -1,10 +1,10 @@
+import cudnn_gbn_lib
+import peer_memory_cuda as pm
 import torch
-from torch.nn.modules.batchnorm import _BatchNorm
-from torch.nn import functional as F
 from torch import Tensor
-import peer_memory_cuda as pm
-import cudnn_gbn_lib
-from torch.cuda.amp import custom_fwd, custom_bwd
+from torch.cuda.amp import custom_bwd, custom_fwd
+from torch.nn import functional as F
+from torch.nn.modules.batchnorm import _BatchNorm
 
 
 class _GroupBatchNorm2d(torch.autograd.Function):
@@ -128,7 +128,7 @@ def __init__(
         affine=True,
         track_running_stats=True,
     ):
-        super(GroupBatchNorm2d, self).__init__(
+        super().__init__(
             num_features,
             eps=eps,
             momentum=momentum,
@@ -165,7 +165,7 @@ def get_peer_buffers(self, num_features):
 
     def _check_input_dim(self, input):
         if input.dim() != 4:
-            raise ValueError("expected 4D input (got {}D input)".format(input.dim()))
+            raise ValueError(f"expected 4D input (got {input.dim()}D input)")
 
     def _check_input_channels(self, input):
         if input.size(1) % 8 != 0:

apex/contrib/examples/gpu_direct_storage/benchmark_load.py

@@ -1,6 +1,8 @@
 import timeit
-import torch
+
 import apex.contrib.gpu_direct_storage as gds
+import torch
+
 
 def run_benchmark_torch_load():
     sizes = [2 ** i for i in range(16, 28)]

apex/contrib/examples/gpu_direct_storage/benchmark_save.py

@@ -1,7 +1,9 @@
 import os
 import timeit
-import torch
+
 import apex.contrib.gpu_direct_storage as gds
+import torch
+
 
 def run_benchmark(func):
     sizes = [2 ** i for i in range(16, 28)]

apex/contrib/examples/gpu_direct_storage/example_load.py

@@ -1,5 +1,5 @@
-import torch
 import apex.contrib.gpu_direct_storage as gds
+import torch
 
 for size in [128, 1024, 8192]:
     x = torch.empty(size, device = "cuda")

apex/contrib/examples/gpu_direct_storage/example_save.py

@@ -1,5 +1,5 @@
-import torch
 import apex.contrib.gpu_direct_storage as gds
+import torch
 
 for size in [128, 1024, 8192]:
     x = torch.linspace(0, 1, size, device = "cuda")

apex/contrib/examples/nccl_allocator/allreduce.py

@@ -1,7 +1,9 @@
 import os
+
 import torch
 import torch.distributed as dist
-import apex.contrib.nccl_allocator as nccl_allocator
+
+from apex.contrib import nccl_allocator
 
 assert os.getenv("WORLD_SIZE") is not None, "Please use: torchrun --nproc-per-node=8 allreduce.py"
 

apex/contrib/examples/nccl_allocator/cache.py

@@ -1,7 +1,9 @@
 import torch
-import apex.contrib.nccl_allocator as nccl_allocator
 from pynvml.smi import nvidia_smi
 
+from apex.contrib import nccl_allocator
+
+
 def set_device(dev):
     import ctypes
     handle = ctypes.CDLL("libcudart.so")

apex/contrib/examples/nccl_allocator/change_cuda_allocator.py

@@ -1,5 +1,6 @@
 import torch
-import apex.contrib.nccl_allocator as nccl_allocator
+
+from apex.contrib import nccl_allocator
 
 nccl_allocator.init()
 nrep = 6

apex/contrib/examples/nccl_allocator/toy_ddp.py

@@ -1,17 +1,17 @@
 import os
+
 import torch
-import torch.nn as nn
-import torch.optim as optim
 import torch.distributed as dist
+from torch import nn, optim
 from torch.nn.parallel import DistributedDataParallel as DDP
 
-import apex.contrib.nccl_allocator as nccl_allocator
+from apex.contrib import nccl_allocator
 
 assert os.getenv("WORLD_SIZE") is not None, "Please use: torchrun --nproc-per-node=8 toy_ddp.py"
 
 class ToyModel(nn.Module):
     def __init__(self):
-        super(ToyModel, self).__init__()
+        super().__init__()
         self.net1 = nn.Linear(10, 10)
         self.relu = nn.ReLU()
         self.net2 = nn.Linear(10, 5)

apex/contrib/focal_loss/__init__.py

@@ -1,6 +1,7 @@
 try:
-    import torch
     import focal_loss_cuda
+    import torch
+
     from .focal_loss import focal_loss
 
     del torch

apex/contrib/focal_loss/focal_loss.py

@@ -1,6 +1,5 @@
-import torch
-
 import focal_loss_cuda
+import torch
 
 
 class FocalLoss(torch.autograd.Function):

apex/contrib/group_norm/group_norm.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-# coding: utf-8
 
 #
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
@@ -8,12 +7,12 @@
 
 import functools
 import os
-import torch
-import torch.nn.init as init
+
 import group_norm_cuda
 import group_norm_v2_cuda
-
+import torch
 from torch import Tensor
+from torch.nn import init
 from torch.nn.parameter import Parameter
 
 __all__ = ["GroupNorm"]

apex/contrib/groupbn/__init__.py

@@ -1,6 +1,7 @@
 try:
-    import torch
     import bnp
+    import torch
+
     from .batch_norm import BatchNorm2d_NHWC
 
     del torch