Torch 2.2 - Part 2 (#979)

.github/workflows/pr-cpu.yaml

@@ -23,6 +23,10 @@ jobs:
           container: mosaicml/pytorch:2.1.0_cpu-python3.10-ubuntu20.04
           markers: "not gpu"
           pytest_command: "coverage run -m pytest"
+        - name: "cpu-2.2.0"
+          container: mosaicml/pytorch:2.2.0_cpu-python3.11-ubuntu20.04
+          markers: "not gpu"
+          pytest_command: "coverage run -m pytest"
     name: ${{ matrix.name }}
     if: github.repository_owner == 'mosaicml'
     with:

.github/workflows/pr-gpu.yaml

@@ -29,6 +29,16 @@ jobs:
           markers: "gpu"
           pytest_command: "coverage run -m pytest"
           deps_group: "all-flash2"
+        - name: "gpu-2.2.0"
+          container: mosaicml/pytorch:2.2.0_cu121-python3.11-ubuntu20.04
+          markers: "gpu"
+          pytest_command: "coverage run -m pytest"
+          deps_group: "all"
+        - name: "gpu-2.2.0-flash2"
+          container: mosaicml/llm-foundry:2.2.0_cu121_flash2-latest
+          markers: "gpu"
+          pytest_command: "coverage run -m pytest"
+          deps_group: "all-flash2"
     name: ${{ matrix.name }}
     if: github.repository_owner == 'mosaicml'
     with:

llmfoundry/optim/__init__.py

@@ -3,9 +3,9 @@
 
 from llmfoundry.optim.adaptive_lion import DecoupledAdaLRLion, DecoupledClipLion
 from llmfoundry.optim.lion import DecoupledLionW
-from llmfoundry.optim.lion8b import DecoupledLionW_8bit
 
 __all__ = [
-    'DecoupledLionW', 'DecoupledLionW_8bit', 'DecoupledClipLion',
-    'DecoupledAdaLRLion'
+    'DecoupledLionW',
+    'DecoupledClipLion',
+    'DecoupledAdaLRLion',
 ]

llmfoundry/utils/builders.py

@@ -37,7 +37,7 @@
                                   ScheduledGarbageCollector)
 from llmfoundry.data.dataloader import build_dataloader
 from llmfoundry.optim import (DecoupledAdaLRLion, DecoupledClipLion,
-                              DecoupledLionW, DecoupledLionW_8bit)
+                              DecoupledLionW)
 from llmfoundry.optim.scheduler import InverseSquareRootWithWarmupScheduler
 from llmfoundry.tokenizers.tiktoken import TiktokenTokenizerWrapper
 
@@ -370,8 +370,6 @@ def build_optimizer(model: torch.nn.Module, name: str,
         return DecoupledClipLion(params, **optimizer_config)
     elif name == 'adalr_lion':
         return DecoupledAdaLRLion(params, **optimizer_config)
-    elif name == 'decoupled_lionw_8b':
-        return DecoupledLionW_8bit(params, **optimizer_config)
     else:
         raise ValueError(f'Not sure how to build optimizer: {name}')
 

setup.py

@@ -104,10 +104,6 @@
     'xentropy-cuda-lib@git+https://github.com/HazyResearch/[email protected]#subdirectory=csrc/xentropy',
 ]
 
-extra_deps['turbo'] = [
-    'mosaicml-turbo==0.0.8',
-]
-
 extra_deps['gpu-flash2'] = [
     'flash-attn==2.5.0',
 ]