Fix FSDP2 breakage in nightly (#2684)

Author: msaroufim

test/dtypes/test_affine_quantized_tensor_parallel.py

@@ -26,6 +26,9 @@
 from torchao.quantization.quant_api import quantize_
 from torchao.utils import TORCH_VERSION_AT_LEAST_2_6
 
+if common_utils.SEED is None:
+    common_utils.SEED = 1234
+
 try:
     import gemlite  # noqa: F401
 

test/dtypes/test_nf4.py

@@ -20,6 +20,7 @@
     apply_activation_checkpointing,
 )
 from torch.distributed.fsdp.wrap import ModuleWrapPolicy
+from torch.testing._internal import common_utils
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_fsdp import FSDPTest
 from torch.testing._internal.common_utils import (
@@ -29,6 +30,9 @@
     run_tests,
 )
 
+if common_utils.SEED is None:
+    common_utils.SEED = 1234
+
 import torchao
 from packaging import version
 from torchao.dtypes._nf4tensor_api import nf4_weight_only

test/prototype/test_quantized_training.py

@@ -15,6 +15,7 @@
 import torch
 import torch.distributed as dist
 import torch.nn.functional as F
+import torch.testing._internal.common_utils as common_utils
 from torch import nn
 from torch.distributed._composable.fsdp import MixedPrecisionPolicy, fully_shard
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
@@ -40,6 +41,9 @@
 )
 from torchao.quantization.quant_api import quantize_
 
+if common_utils.SEED is None:
+    common_utils.SEED = 1234
+
 _DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])
 
 

test/test_low_bit_optim.py

@@ -16,6 +16,7 @@
     OffloadPolicy,
     fully_shard,
 )
+from torch.testing._internal import common_utils
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_fsdp import FSDPTest
 from torch.testing._internal.common_utils import (
@@ -25,6 +26,9 @@
     run_tests,
 )
 
+if common_utils.SEED is None:
+    common_utils.SEED = 1234
+
 from packaging.version import Version
 from torchao import optim
 from torchao.optim.quant_utils import (