Move f/utils => f/internal; move testing libs to fair_dev/testing (#1004

)

benchmarks/experimental/experimental_async_approaches.py

@@ -21,12 +21,12 @@
 import torchtext
 from torchtext.data.utils import get_tokenizer
 
+from fair_dev.testing.testing import dist_init, get_worker_map
 from fairscale.experimental.nn.ampnet_pipe import pipe
 from fairscale.nn.model_parallel import initialize_model_parallel
 from fairscale.nn.model_parallel.initialize import get_pipeline_parallel_group
 from fairscale.nn.pipe import LazyModule
 from fairscale.optim import GradScaler
-from fairscale.utils.testing import dist_init, get_worker_map
 
 try:
     from fairscale.optim import Adam  # type: ignore

benchmarks/pipe.py

@@ -16,9 +16,9 @@
 import utils
 
 from benchmarks.golden_configs.lm_wikitext2 import Pipe as lm_wikitext2
+from fair_dev.testing.testing import dist_init
 from fairscale.nn import Pipe
 from fairscale.nn.model_parallel import initialize_model_parallel
-from fairscale.utils.testing import dist_init
 
 MPI_PORT = 29500
 RPC_PORT = 29501

fairscale/utils/testing.py → fair_dev/testing/testing.py

@@ -49,9 +49,9 @@
 import torch.multiprocessing as mp
 import torch.nn as nn
 
+from fairscale.internal import torch_version
 from fairscale.nn.model_parallel import destroy_model_parallel, initialize_model_parallel
 from fairscale.nn.model_parallel.random import model_parallel_cuda_manual_seed
-from fairscale.utils import torch_version
 
 if TYPE_CHECKING:
     Base = nn.Module[Tensor]

fairscale/experimental/nn/distributed_pipeline/pipeline.py

@@ -10,8 +10,8 @@
 from torch import Tensor, nn
 from torch.distributed import rpc
 
+from fairscale.internal import torch_version
 from fairscale.nn.pipe import microbatch
-from fairscale.utils import torch_version
 
 from .data import DataConsumer
 from .graph import Node, PipelineModulesGraph

fairscale/experimental/nn/ssd_offload.py

@@ -17,7 +17,7 @@
 import torch
 from torch.serialization import DEFAULT_PROTOCOL as DEFAULT_PROTOCOL
 
-from fairscale.utils import torch_version
+from fairscale.internal import torch_version
 
 try:
     from torch.utils._pytree import tree_map

fairscale/experimental/nn/sync_batchnorm.py

@@ -10,8 +10,8 @@
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
 
+from fairscale.internal import torch_version
 from fairscale.nn.checkpoint import is_checkpointing, is_recomputing
-from fairscale.utils import torch_version
 
 
 def _forward(input: Tensor, affine: bool, mean: Tensor, invstd: Tensor, weight: Tensor, bias: Tensor) -> Tensor:

fairscale/nn/checkpoint/checkpoint_activations.py

@@ -14,7 +14,7 @@
 import torch.nn as nn
 import torch.utils.checkpoint as torch_checkpoint
 
-from fairscale.utils.containers import pack_kwargs, split_non_tensors, unpack_kwargs, unpack_non_tensors
+from fairscale.internal.containers import pack_kwargs, split_non_tensors, unpack_kwargs, unpack_non_tensors
 
 from .checkpoint_utils import patch_batchnorm
 

fairscale/nn/data_parallel/fully_sharded_data_parallel.py

@@ -40,19 +40,19 @@
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
 
-from fairscale.nn.misc import FlattenParamsWrapper
-from fairscale.nn.wrap import auto_wrap, config_auto_wrap_policy, enable_wrap
-from fairscale.utils.containers import apply_to_tensors
-from fairscale.utils.parallel import (
+from fairscale.internal.containers import apply_to_tensors
+from fairscale.internal.parallel import (
     ProcessGroupName,
     chunk_and_pad,
     enable_pytorch_sync_bn,
     get_process_group_cached,
     validate_process_group,
 )
-from fairscale.utils.params import calc_grad_norm, recursive_copy_to_device
-from fairscale.utils.reduce_scatter_bucketer import ReduceScatterBucketer
-from fairscale.utils.state_dict import replace_by_prefix_
+from fairscale.internal.params import calc_grad_norm, recursive_copy_to_device
+from fairscale.internal.reduce_scatter_bucketer import ReduceScatterBucketer
+from fairscale.internal.state_dict import replace_by_prefix_
+from fairscale.nn.misc import FlattenParamsWrapper
+from fairscale.nn.wrap import auto_wrap, config_auto_wrap_policy, enable_wrap
 
 from . import fsdp_optim_utils as ou
 

fairscale/nn/data_parallel/sharded_ddp.py

@@ -21,9 +21,9 @@
 import torch.autograd.profiler as profiler
 import torch.distributed as dist
 
+from fairscale.internal.params import Workhandle, get_global_rank
 from fairscale.nn.misc import GradBucket
 from fairscale.optim import OSS
-from fairscale.utils.params import Workhandle, get_global_rank
 
 
 def _trainable(param: torch.Tensor) -> bool:

fairscale/nn/misc/flatten_params_wrapper.py

@@ -44,7 +44,7 @@
     import_ssd_offload = False
     pass
 
-from fairscale.utils.state_dict import replace_by_prefix_
+from fairscale.internal.state_dict import replace_by_prefix_
 
 if TYPE_CHECKING:
     from collections import OrderedDict  # noqa: F401

fairscale/nn/pipe/messages.py

@@ -11,8 +11,8 @@
 
 import torch
 
+from fairscale.internal.object import pyobject_to_tensor, tensor_to_pyobject
 from fairscale.nn.model_parallel import get_pipeline_parallel_group
-from fairscale.utils.object import pyobject_to_tensor, tensor_to_pyobject
 
 from .types import MESSAGE_GENERATION_START, InputDevice, PipeMessage, Tensors
 

fairscale/nn/pipe/pipe.py

@@ -27,7 +27,7 @@
 import torch.autograd
 import torch.cuda
 
-from fairscale.utils import torch_version
+from fairscale.internal import torch_version
 
 from . import microbatch
 from .batchnorm import DeferredBatchNorm

fairscale/optim/grad_scaler.py

@@ -18,7 +18,7 @@
 from torch.optim import Optimizer
 from torch.optim.sgd import SGD
 
-from fairscale.utils import torch_version
+from fairscale.internal import torch_version
 
 
 class _GeneralMultiDeviceReplicator(object):

fairscale/optim/oss.py

@@ -17,8 +17,8 @@
 from torch.nn import Parameter
 from torch.optim import SGD, Optimizer
 
+from fairscale.internal.params import calc_grad_norm, get_global_rank, recursive_copy_to_device
 from fairscale.nn.misc import ParamBucket
-from fairscale.utils.params import calc_grad_norm, get_global_rank, recursive_copy_to_device
 
 __all__ = ["OSS"]
 

tests/experimental/nn/ampnet_pipe_process/test_ampnet_pipe.py

@@ -22,8 +22,8 @@
 from torch.optim.optimizer import Optimizer
 from torch.utils.data import DataLoader, Dataset
 
+from fair_dev.testing.testing import get_worker_map, torch_spawn
 from fairscale.experimental.nn.ampnet_pipe.pipe import AMPnetPipe
-from fairscale.utils.testing import get_worker_map, torch_spawn
 
 
 class MySGD(Optimizer):

tests/experimental/nn/data_parallel/test_gossip.py

@@ -15,8 +15,8 @@
 import torch.distributed
 import torch.nn.functional as F
 
+from fair_dev.testing.testing import skip_if_single_gpu, spawn_for_all_world_sizes
 import fairscale.experimental.nn.data_parallel.gossip as gossip
-from fairscale.utils.testing import skip_if_single_gpu, spawn_for_all_world_sizes
 
 # Enfore CUBLAS reproducibility, see https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
 os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

tests/experimental/nn/test_auto_shard.py

@@ -14,7 +14,7 @@
 import torch.nn
 import torch.nn as nn
 
-from fairscale.utils import torch_version
+from fairscale.internal import torch_version
 
 
 class PositionalEncoding(nn.Module):

tests/experimental/nn/test_mevo.py

@@ -12,9 +12,9 @@
 import pytest
 import torch
 
+from fair_dev.testing.testing import skip_if_no_cuda
 from fairscale.experimental.nn import MEVO
 from fairscale.experimental.nn.mevo import BaselineSoftmaxNllLoss, get_data
-from fairscale.utils.testing import skip_if_no_cuda
 
 
 @pytest.fixture(scope="session", params=[torch.float16, torch.float32])

tests/experimental/nn/test_multiprocess_pipe.py

@@ -20,9 +20,9 @@
 import torch.multiprocessing as mp
 import torch.nn as nn
 
+from fair_dev.testing.testing import skip_if_single_gpu
 from fairscale.experimental.nn.distributed_pipeline import DistributedLoss, DistributedPipeline, PipelineModulesGraph
-from fairscale.utils import torch_version
-from fairscale.utils.testing import skip_if_single_gpu
+from fairscale.internal import torch_version
 
 pytestmark = pytest.mark.skipif(
     not torch.cuda.is_available() or torch_version() < (1, 9, 0),

tests/experimental/nn/test_offload.py

@@ -14,9 +14,9 @@
 import pytest
 import torch
 
+from fair_dev.testing.testing import skip_if_no_cuda
 from fairscale.experimental.nn.offload import OffloadModel
-from fairscale.utils import torch_version
-from fairscale.utils.testing import skip_if_no_cuda
+from fairscale.internal import torch_version
 
 if torch_version() >= (1, 8, 0):
     from fairscale.experimental.nn.auto_shard import shard_model

tests/experimental/tooling/test_layer_memory_tracker.py

@@ -10,13 +10,13 @@
 import torch.nn as nn
 from torch.nn.parallel import DistributedDataParallel
 
+from fair_dev.testing.testing import GPT2, dist_init, skip_if_no_cuda, skip_if_single_gpu, temp_files_ctx
 from fairscale.experimental.tooling.layer_memory_tracker import (
     LayerwiseMemoryTracker,
     ProcessGroupTracker,
     find_best_reset_points,
 )
 from fairscale.nn import FullyShardedDataParallel
-from fairscale.utils.testing import GPT2, dist_init, skip_if_no_cuda, skip_if_single_gpu, temp_files_ctx
 
 
 @skip_if_no_cuda()

tests/nn/checkpoint/test_checkpoint_activations.py

@@ -10,11 +10,11 @@
 import torch.nn as nn
 from torch.utils.checkpoint import checkpoint as torch_checkpoint_wrapper
 
+from fair_dev.testing.testing import skip_if_no_cuda
+from fairscale.internal import torch_version
 from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper, disable_checkpointing
 from fairscale.nn.misc import FlattenParamsWrapper
 from fairscale.nn.misc import checkpoint_wrapper as deprecated_checkpoint_wrapper
-from fairscale.utils import torch_version
-from fairscale.utils.testing import skip_if_no_cuda
 
 
 def get_cuda_mem_allocated():

tests/nn/checkpoint/test_checkpoint_activations_norm.py

@@ -14,9 +14,9 @@
 from torch.nn import BatchNorm2d, LayerNorm, Linear, Sequential
 from torch.optim import SGD
 
+from fair_dev.testing.testing import objects_are_equal
+from fairscale.internal import torch_version
 from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
-from fairscale.utils import torch_version
-from fairscale.utils.testing import objects_are_equal
 
 NORM_TYPES = [LayerNorm, BatchNorm2d]
 MP_TYPES = ["fp32", "fp16", "call_half"]

tests/nn/data_parallel/test_fsdp.py

@@ -18,10 +18,7 @@
 from torch import nn
 import torch.distributed
 
-from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
-from fairscale.nn.data_parallel import FullyShardedDataParallel, TrainingState
-from fairscale.utils import torch_version
-from fairscale.utils.testing import (
+from fair_dev.testing.testing import (
     DeviceAndTypeCheckModule,
     DummyProcessGroup,
     dist_init,
@@ -30,6 +27,9 @@
     skip_a_test_if_in_CI,
     spawn_for_all_world_sizes,
 )
+from fairscale.internal import torch_version
+from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
+from fairscale.nn.data_parallel import FullyShardedDataParallel, TrainingState
 
 if torch_version() >= (1, 8, 0):
     from fairscale.optim.grad_scaler import ShardedGradScaler

tests/nn/data_parallel/test_fsdp_apply.py

@@ -10,7 +10,7 @@
 import pytest
 import torch.nn as nn
 
-from fairscale.utils import torch_version
+from fairscale.internal import torch_version
 
 from .test_fsdp import (
     CONFIG_OPTIONS,

tests/nn/data_parallel/test_fsdp_freezing_weights.py

@@ -21,8 +21,8 @@
 from torch.nn.parallel import DistributedDataParallel
 import torch.optim as optim
 
+from fair_dev.testing.testing import dist_init, objects_are_equal, rmf, skip_if_single_gpu, teardown
 from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
-from fairscale.utils.testing import dist_init, objects_are_equal, rmf, skip_if_single_gpu, teardown
 
 
 class FreezeModel(nn.Module):

tests/nn/data_parallel/test_fsdp_grad_acc.py

@@ -12,8 +12,8 @@
 from parameterized import parameterized
 import torch
 
+from fair_dev.testing.testing import DummyProcessGroup, make_cudnn_deterministic, objects_are_equal
 from fairscale.nn.data_parallel import FullyShardedDataParallel
-from fairscale.utils.testing import DummyProcessGroup, make_cudnn_deterministic, objects_are_equal
 
 from .test_fsdp import DistributedTest, NestedWrappedModule, rename_test, spawn_and_init
 

tests/nn/data_parallel/test_fsdp_hf_transformer_eval.py

@@ -6,9 +6,9 @@
 import torch
 from torch import nn
 
+from fair_dev.testing.testing import dist_init
 from fairscale.nn import FullyShardedDataParallel as FSDP
 from fairscale.nn import auto_wrap, enable_wrap
-from fairscale.utils.testing import dist_init
 
 
 def wrap_transformer_only(module, recurse, **kwargs):

tests/nn/data_parallel/test_fsdp_input.py

@@ -16,10 +16,10 @@
 from torch.nn import Linear, Module
 from torch.optim import SGD
 
+from fair_dev.testing.testing import dist_init, rmf, skip_if_no_cuda, teardown
+from fairscale.internal import torch_version
 from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
 from fairscale.nn.data_parallel import TrainingState
-from fairscale.utils import torch_version
-from fairscale.utils.testing import dist_init, rmf, skip_if_no_cuda, teardown
 
 
 # A fixture to get tempfiles and ensure they are cleaned up.

tests/nn/data_parallel/test_fsdp_memory.py

@@ -18,12 +18,12 @@
 from torch.nn.parallel import DistributedDataParallel
 import torch.optim as optim
 
+from fair_dev.testing.testing import dist_init, dump_all_tensors, skip_if_single_gpu, teardown, temp_files_ctx
+from fairscale.internal import torch_version
+from fairscale.internal.parallel import get_process_group_cached
 from fairscale.nn import checkpoint_wrapper
 from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
 from fairscale.nn.data_parallel import auto_wrap_bn
-from fairscale.utils import torch_version
-from fairscale.utils.parallel import get_process_group_cached
-from fairscale.utils.testing import dist_init, dump_all_tensors, skip_if_single_gpu, teardown, temp_files_ctx
 
 
 def to_fsdp(module, fsdp_config):

tests/nn/data_parallel/test_fsdp_metadata.py

@@ -14,8 +14,8 @@
 import torch.nn as nn
 from torch.optim import Adam
 
+from fair_dev.testing.testing import in_temporary_directory, skip_if_single_gpu, temp_files_ctx
 from fairscale.nn import FullyShardedDataParallel
-from fairscale.utils.testing import in_temporary_directory, skip_if_single_gpu, temp_files_ctx
 from tests.nn.data_parallel.test_fsdp import DistributedTest, MixtureOfExperts, rename_test, spawn_and_init
 
 USE_TEMPFILE = True  # False for debugging

tests/nn/data_parallel/test_fsdp_multiple_forward.py

@@ -17,10 +17,10 @@
 from torch.nn import Linear, Module
 from torch.optim import SGD
 
+from fair_dev.testing.testing import dist_init, skip_if_single_gpu, teardown
+from fairscale.internal import torch_version
 from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
 from fairscale.nn.data_parallel import TrainingState
-from fairscale.utils import torch_version
-from fairscale.utils.testing import dist_init, skip_if_single_gpu, teardown
 
 
 def _test_func(rank, world_size, fsdp_config, tempfile_name, unused):

tests/nn/data_parallel/test_fsdp_multiple_forward_checkpoint.py

@@ -20,12 +20,12 @@
 from torch.nn.parallel import DistributedDataParallel
 import torch.optim as optim
 
+from fair_dev.testing.testing import dist_init, skip_if_single_gpu, teardown, temp_files_ctx
+from fairscale.internal import torch_version
 from fairscale.nn import checkpoint_wrapper
 from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
 from fairscale.nn.data_parallel import auto_wrap_bn
 from fairscale.nn.wrap import enable_wrap, wrap
-from fairscale.utils import torch_version
-from fairscale.utils.testing import dist_init, skip_if_single_gpu, teardown, temp_files_ctx
 
 
 class Model(nn.Module):