Update FSDP checkpointing test to use UC Volumes and updated dockerfile for new composer version (#3865)

Author: rithwik-db

docker/README.md

@@ -15,8 +15,8 @@ all dependencies for both NLP and Vision models. They are built on top of the
 <!-- BEGIN_COMPOSER_BUILD_MATRIX -->
 | Composer Version   | CUDA Support   | Docker Tag                                                     |
 |--------------------|----------------|----------------------------------------------------------------|
-| 0.30.0             | Yes            | `mosaicml/composer:latest`, `mosaicml/composer:0.30.0`         |
-| 0.30.0             | No             | `mosaicml/composer:latest_cpu`, `mosaicml/composer:0.30.0_cpu` |
+| 0.31.0             | Yes            | `mosaicml/composer:latest`, `mosaicml/composer:0.31.0`         |
+| 0.31.0             | No             | `mosaicml/composer:latest_cpu`, `mosaicml/composer:0.31.0_cpu` |
 <!-- END_COMPOSER_BUILD_MATRIX -->
 
 **Note**: For a lightweight installation, we recommended using a [MosaicML PyTorch Image](#pytorch-images) and manually

docker/build_matrix.yaml

@@ -183,35 +183,35 @@
   TORCHVISION_VERSION: 0.21.0
   UBUNTU_VERSION: '20.04'
 - BASE_IMAGE: nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04
-  COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.30.0
+  COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.31.0
   CUDA_VERSION: 12.6.3
   EFA_INSTALLER_VERSION: ''
-  IMAGE_NAME: composer-0-30-0
+  IMAGE_NAME: composer-0-31-0
   MOFED_VERSION: latest-23.10
   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.12'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.7.0
   TAGS:
-  - mosaicml/composer:0.30.0
+  - mosaicml/composer:0.31.0
   - mosaicml/composer:latest
   TARGET: composer_stage
   TORCHVISION_VERSION: 0.22.0
   UBUNTU_VERSION: '22.04'
 - BASE_IMAGE: ubuntu:22.04
-  COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.30.0
+  COMPOSER_INSTALL_COMMAND: mosaicml[all]==0.31.0
   CUDA_VERSION: ''
   EFA_INSTALLER_VERSION: ''
-  IMAGE_NAME: composer-0-30-0-cpu
+  IMAGE_NAME: composer-0-31-0-cpu
   MOFED_VERSION: latest-23.10
   NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.12'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
   PYTORCH_VERSION: 2.7.0
   TAGS:
-  - mosaicml/composer:0.30.0_cpu
+  - mosaicml/composer:0.31.0_cpu
   - mosaicml/composer:latest_cpu
   TARGET: composer_stage
   TORCHVISION_VERSION: 0.22.0

docker/generate_build_matrix.py

@@ -294,7 +294,7 @@ def _main():
     composer_entries = []
 
     # The `GIT_COMMIT` is a placeholder and Jenkins will substitute it with the actual git commit for the `composer_staging` images
-    composer_versions = ['0.30.0']  # Only build images for the latest composer version
+    composer_versions = ['0.31.0']  # Only build images for the latest composer version
     composer_python_versions = [PRODUCTION_PYTHON_VERSION]  # just build composer against the latest
 
     for product in itertools.product(composer_python_versions, composer_versions, cuda_options):

tests/fixtures/fixtures.py

@@ -131,6 +131,14 @@ def s3_read_only_prefix():
     return 'read_only'
 
 
+@pytest.fixture
+def uc_volume_path(request: pytest.FixtureRequest):
+    if request.node.get_closest_marker('remote') is None:
+        return 'my-volume'
+    else:
+        return os.environ.get('UC_VOLUME_PATH', 'Volumes/main/regression_testing/composer_artifacts/')
+
+
 ## MODEL HELPERS ##
 def causal_lm_model_helper(config):  # type: ignore
     transformers = pytest.importorskip('transformers')

tests/trainer/test_fsdp_checkpoint.py

@@ -20,6 +20,7 @@
 from packaging import version
 from torch.distributed._shard.sharded_tensor import ShardedTensor
 from torch.distributed._tensor import DTensor
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.utils.data import DataLoader
 from torchmetrics import Metric, MetricCollection
 from torchmetrics.classification import MulticlassAccuracy
@@ -33,7 +34,7 @@
 from composer.utils import FSDPConfig, TPConfig, dist, parse_uri
 from composer.utils.checkpoint import dist_cp_load
 from composer.utils.file_helpers import get_file
-from composer.utils.object_store import S3ObjectStore
+from composer.utils.object_store import UCObjectStore
 from composer.utils.reproducibility import get_rng_state
 from tests.common import RandomClassificationDataset, deep_compare
 from tests.common.markers import world_size
@@ -517,6 +518,7 @@ def test_fsdp_mixed_with_sync(
         '0.28.0',
         '0.29.0',
         '0.30.0',
+        '0.31.0',
     ],
 )
 @pytest.mark.filterwarnings(r'ignore:.*metrics are not saved with sharded state dict.*:UserWarning')
@@ -529,8 +531,7 @@ def test_fsdp_load_old_checkpoint(
     precision: str,
     sharding_strategy: str,
     state_dict_type: str,
-    s3_bucket: str,
-    s3_read_only_prefix: str,
+    uc_volume_path: str,
     composer_version: str,
 ):
     if composer_version == '0.18.1' and state_dict_type == 'full' and precision == 'amp_bf16' and sharding_strategy == 'FULL_SHARD':
@@ -540,25 +541,27 @@ def test_fsdp_load_old_checkpoint(
         if state_dict_type == 'sharded':
             pytest.skip('Loading legacy sharded checkpoints are not supported after v0.25.0.')
 
-        load_path_dir = (
-            f's3://{s3_bucket}/{s3_read_only_prefix}/backwards_compatibility/'
-            f'{composer_version}/{sharding_strategy.lower()}_{state_dict_type}_'
-            f'{precision}/'
+        load_path_dir = os.path.join(
+            f'dbfs:/{uc_volume_path}',
+            'backwards_compatibility',
+            composer_version,
+            f'{sharding_strategy.lower()}_{state_dict_type}_{precision}',
         )
         if ((version.parse(composer_version) > version.parse('0.15.0')) and state_dict_type != 'full'):
-            load_path_dir = (load_path_dir + 'ep0-ba2/')
+            load_path_dir = os.path.join(load_path_dir, 'ep0-ba2')
 
-        load_path = load_path_dir + f'ba2_rank0.pt'
+        load_path = os.path.join(load_path_dir, f'ba2_rank0.pt')
     else:
-        load_path = (
-            f's3://{s3_bucket}/{s3_read_only_prefix}/backwards_compatibility/'
-            f'{composer_version}/{sharding_strategy.lower()}_{state_dict_type}_'
-            f'{precision}/'
+        load_path = os.path.join(
+            f'dbfs:/{uc_volume_path}',
+            'backwards_compatibility',
+            composer_version,
+            f'{sharding_strategy.lower()}_{state_dict_type}_{precision}',
         )
         if state_dict_type == 'full':
-            load_path += 'ba2_rank0.pt'
+            load_path = os.path.join(load_path, 'ba2_rank0.pt')
         else:
-            load_path += 'ep0-ba2/'
+            load_path = os.path.join(load_path, 'ep0-ba2')
 
     if composer_version == '0.15.1':
         num_classes = 8  # This parameter setting is very important. Don't change or the test will fail.
@@ -619,7 +622,7 @@ def test_fsdp_load_old_checkpoint(
                 'rng': get_rng_state(),
             }
 
-            object_store = S3ObjectStore(bucket=f'{s3_bucket}')
+            object_store = UCObjectStore(path=uc_volume_path)
             storage_reader = DistCPObjectStoreReader(
                 source_path=parsed_load_path,
                 destination_path=destination,