Train with DeepSpeed ZeRO-2 and ZeRO-3 (#13587)

* Train with DeepSpeed ZeRO-2 and ZeRO-3

- Single device: keep Accelerator.distributed_type as DEEPSPEED. Reporting NO
  turned on accelerate's native AMP, whose DeepSpeed fallthrough is a
  dtype-less (float16) autocast over a bfloat16 model: Half vs BFloat16 in
  the fused LoRA kernel on the first step.
- ZeRO-3 init: drop the device_map Unsloth picked; transformers refuses any.
- ZeRO-3 training: Unsloth's fast paths read child weights (norms, fused
  LoRA bases, lm_head in the fused loss) from the parent's forward, which
  ZeRO-3 does not gather. Make each repeated block a ZeRO-3 leaf module and
  register every other weight as an external parameter of its parent.

Fixes #225

* ZeRO-3 leaves: only the outermost repeated blocks

A one-layer model was missed by the len > 1 rule, and an MoE experts list
inside a decoder layer became a leaf of its own, whose routing-dependent
all-gathers differ across ranks. Walk top-down, accept a one-element list,
and never look inside a chosen block.

* ZeRO-3: rely on DeepSpeed's own gather, keep a caller's device_map

The leaf / external-parameter setup never ran (deepspeed.zero is not an
importable module, and the broad except hid it) and is not needed: during
the engine forward DeepSpeed gathers a partitioned weight on any eager read
and registers it as external. Removed with its tests.

Under ZeRO-3 init only a placement left to Unsloth is dropped; a device_map
the caller chose reaches transformers and raises as before.

* Tests: stub deepspeed_zero3_device_map in the loader fixture, drop unused import

tests/test_accelerator_distributed_type_patch.py

@@ -15,15 +15,16 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 """`_utils.py` must patch `Accelerator.distributed_type` with a property, not a bare
 function: a function binds as a method and inverts accelerate's `!= NO` device_map
-guards on a single device (#10016). `import unsloth` needs CUDA or XPU, so these run
-the statement from source in a fresh interpreter, which also catches a revert.
+guards on a single device (#10016). It must still read as DEEPSPEED under DeepSpeed, or
+accelerate turns on a float16 autocast over a bfloat16 model (#225). `import unsloth` needs
+CUDA or XPU, so these run the statements from source in a fresh interpreter, which also
+catches a revert.
 """
 
 import ast
 import pathlib
 import subprocess
 import sys
-import textwrap
 
 import pytest
 
@@ -47,6 +48,43 @@ def _patch_statement() -> ast.Assign:
     return matches[0]
 
 
+def _getter_source() -> str:
+    """The function the property wraps, so the fresh interpreter can define it first."""
+    tree = ast.parse(_UTILS.read_text(encoding = "utf-8"))
+    name = ast.unparse(_patch_statement().value.args[0])
+    defs = [n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == name]
+    assert len(defs) == 1, f"expected one {name}, found {len(defs)}"
+    return ast.unparse(defs[0])
+
+
+def _read_distributed_type(state_type):
+    """repr of the patched attribute and its `== NO` / `!= NO` on an Accelerator whose state says `state_type`."""
+    script = "\n".join(
+        [
+            "import types",
+            "import accelerate.accelerator",
+            "from accelerate.utils.dataclasses import DistributedType",
+            _getter_source(),
+            ast.unparse(_patch_statement()),
+            # Read it off an instance without standing up a real distributed backend.
+            "accelerator = accelerate.accelerator.Accelerator.__new__(accelerate.accelerator.Accelerator)",
+            *(
+                []
+                if state_type is None
+                else [
+                    f"accelerator.state = types.SimpleNamespace(distributed_type = DistributedType.{state_type})"
+                ]
+            ),
+            "value = accelerator.distributed_type",
+            "print(repr(value))",
+            "print(value == DistributedType.NO, value != DistributedType.NO)",
+        ]
+    )
+    result = subprocess.run([sys.executable, "-c", script], capture_output = True, text = True)
+    assert result.returncode == 0, result.stderr
+    return result.stdout.strip().splitlines()[-2:]
+
+
 def test_distributed_type_is_replaced_with_a_property():
     value = _patch_statement().value
     assert isinstance(value, ast.Call) and ast.unparse(value.func) == "property", (
@@ -57,26 +95,17 @@ def test_distributed_type_is_replaced_with_a_property():
 
 def test_patched_distributed_type_reads_as_no():
     """Running the real assignment must leave the attribute reading as NO."""
-    statement = ast.unparse(_patch_statement())
-    script = textwrap.dedent(
-        f"""
-        import accelerate.accelerator
-        from accelerate.utils.dataclasses import DistributedType
+    DistributedType = accelerate.utils.dataclasses.DistributedType
+    for state_type in (None, "NO", "MULTI_GPU", "FSDP"):
+        value, comparisons = _read_distributed_type(state_type)
+        assert value == repr(DistributedType.NO), (state_type, value)
+        # The second is what accelerate's device_map guards actually test.
+        assert comparisons == "True False", (state_type, comparisons)
 
-        {statement}
 
-        # Read it off an instance without standing up a real distributed backend.
-        accelerator = accelerate.accelerator.Accelerator.__new__(
-            accelerate.accelerator.Accelerator
-        )
-        value = accelerator.distributed_type
-        print(repr(value))
-        print(value == DistributedType.NO, value != DistributedType.NO)
-        """
-    )
-    result = subprocess.run([sys.executable, "-c", script], capture_output = True, text = True)
-    assert result.returncode == 0, result.stderr
-    value, comparisons = result.stdout.strip().splitlines()[-2:]
-    assert value == repr(accelerate.utils.dataclasses.DistributedType.NO), value
-    # The second is what accelerate's device_map guards actually test.
-    assert comparisons == "True False", comparisons
+def test_patched_distributed_type_keeps_deepspeed():
+    """Under DeepSpeed, NO enables native AMP and get_mixed_precision_context_manager's
+    DeepSpeed fallthrough is a dtype-less, i.e. float16, autocast."""
+    value, comparisons = _read_distributed_type("DEEPSPEED")
+    assert value == repr(accelerate.utils.dataclasses.DistributedType.DEEPSPEED), value
+    assert comparisons == "False True", comparisons

tests/test_deepspeed_zero3.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""DeepSpeed ZeRO-3 (#225): the load must not pass a device_map transformers refuses."""
+
+import pytest
+
+transformers = pytest.importorskip("transformers")
+
+from unsloth.models import loader_utils
+
+
+@pytest.fixture
+def zero3(monkeypatch):
+    def set_enabled(enabled):
+        monkeypatch.setattr(
+            transformers.integrations, "is_deepspeed_zero3_enabled", lambda: enabled
+        )
+
+    return set_enabled
+
+
+def test_zero3_drops_the_placement_unsloth_picked(zero3):
+    zero3(True)
+    for device_map in (loader_utils.DEFAULT_DEVICE_MAP, "sequential", "auto", "cuda"):
+        assert loader_utils.deepspeed_zero3_device_map(device_map) is None, device_map
+
+
+@pytest.mark.parametrize(
+    "device_map", [{"model.layers.0": 0, "lm_head": 1}, {"": "cuda:1"}, "cuda:1", None]
+)
+def test_zero3_keeps_a_map_the_caller_chose(zero3, device_map):
+    """transformers rejects these under ZeRO-3; silently dropping one would move the run to another card."""
+    zero3(True)
+    assert loader_utils.deepspeed_zero3_device_map(device_map) is device_map
+
+
+def test_without_zero3_the_device_map_is_untouched(zero3):
+    zero3(False)
+    for device_map in (loader_utils.DEFAULT_DEVICE_MAP, "sequential", {"": "cuda:1"}, None):
+        assert loader_utils.deepspeed_zero3_device_map(device_map) is device_map

tests/test_optimized_precision_conflicts.py

@@ -66,6 +66,7 @@ def loader():
         _requested_float32 = lambda dtype: False,
         hf_login = lambda token: token,
         requested_device_map = lambda device: device,
+        deepspeed_zero3_device_map = lambda device: device,
         is_automatic_device_map = lambda device: isinstance(device, str),
         prepare_device_map = lambda: ("sequential", False),
         ALLOW_BITSANDBYTES = True,

unsloth/models/_utils.py

@@ -4356,8 +4356,17 @@ if DEVICE_COUNT == 1 and int(os.environ.get("WORLD_SIZE", "1")) <= 1:
     import accelerate.state
 
     accelerate.state.PartialState._prepare_backend = _prepare_backend
+
+    # DeepSpeed keeps its type even on one device: reporting NO turns on accelerate's native AMP, whose
+    # dtype-less autocast under DeepSpeed is float16 and breaks a bfloat16 model (#225).
+    def _single_device_distributed_type(self):
+        real = getattr(self.__dict__.get("state"), "distributed_type", None)
+        return (
+            DistributedType.DEEPSPEED if real == DistributedType.DEEPSPEED else DistributedType.NO
+        )
+
     # Must be a property: a bare function binds as a method, inverting every `!= DistributedType.NO` guard in accelerate (#10016).
-    accelerate.accelerator.Accelerator.distributed_type = property(lambda self: DistributedType.NO)
+    accelerate.accelerator.Accelerator.distributed_type = property(_single_device_distributed_type)
 
 
 _PER_LAYER_DEVICE_MISSING = object()

unsloth/models/loader.py

@@ -45,6 +45,7 @@ from .lora_init import adapter_used_fast_pissa, fast_lora_init, record_fast_piss
 from ..tokenizers_v1 import tokenizers_v1_on_return
 from .loader_utils import (
     DEFAULT_DEVICE_MAP,
+    deepspeed_zero3_device_map,
     OFFLOAD_EMBEDDING_AUTO,
     _exclude_rope_inv_freq_from_ddp,
     _get_fp8_mode_and_check_settings,
@@ -1114,7 +1115,7 @@ class FastLanguageModel(FastLlamaModel):
                 dtype = bnb_compute_dtype
 
         # Distributed-safe placement: under torchrun / accelerate launch each rank must load on its own device. Quantized weights make it mandatory, since Accelerate raises device relocation errors on them, but a 16-bit load needs it just as much: an unchosen placement ends up dispatching to cuda:0, so all the ranks land on card 0 and contend for it while the rest of the cards stay empty (#3459). A device the caller named is left alone.
-        device_map = requested_device_map(device_map)
+        device_map = deepspeed_zero3_device_map(requested_device_map(device_map))
         if is_automatic_device_map(device_map):
             distributed_device_map, is_dist = prepare_device_map()
             if is_dist:
@@ -2096,7 +2097,7 @@ class FastModel(FastBaseModel):
             qat_scheme = "int8-int4"
 
         # Distributed-safe placement: under torchrun / accelerate launch each rank must load on its own device. Quantized weights make it mandatory, since Accelerate raises device relocation errors on them, but a 16-bit load needs it just as much: an unchosen placement ends up dispatching to cuda:0, so all the ranks land on card 0 and contend for it while the rest of the cards stay empty (#3459). A device the caller named is left alone.
-        device_map = requested_device_map(device_map)
+        device_map = deepspeed_zero3_device_map(requested_device_map(device_map))
         if is_automatic_device_map(device_map):
             distributed_device_map, is_dist = prepare_device_map()
             if is_dist:

unsloth/models/loader_utils.py

@@ -197,6 +197,19 @@ def prepare_device_map():
     return device_map, True
 
 
+def deepspeed_zero3_device_map(device_map):
+    """None under DeepSpeed ZeRO-3 init: transformers refuses any device_map there, since ZeRO-3 places the
+    weights itself (#225). Only a placement left to us is dropped; a map the caller chose still raises."""
+    try:
+        from transformers.integrations import is_deepspeed_zero3_enabled
+        zero3 = is_deepspeed_zero3_enabled()
+    except Exception:
+        zero3 = False
+    if not zero3:
+        return device_map
+    return None if is_automatic_device_map(device_map) else device_map
+
+
 UNSLOTH_DEVICE_MAP = "unsloth"
 
 # Same planner, different answer when it declines: every veto path ends in "sequential", which fills the first device to its whole free budget, wrong for a caller that picked several cards.