muellerzr/github-pr-history huggingface dataset:
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 9ee840440..599e710f0 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -752,7 +752,7 @@ def launch_command(args):
     # Get the default from the config file.
     if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu:
         defaults = load_config_from_file(args.config_file)
-        if not args.multi_gpu and not args.tpu and not args.use_deepspeed:
+        if not args.multi_gpu and not args.tpu and not args.use_deepspeed and not args.use_fsdp:
             args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index bd4b72682..7100d2c88 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -127,7 +127,7 @@ def set_module_tensor_to_device(
 
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif torch.device(device) != module._parameters[tensor_name].device:
+        elif value is not None or torch.device(device) != module._parameters[tensor_name].device:
             param_cls = type(module._parameters[tensor_name])
             kwargs = module._parameters[tensor_name].__dict__
             new_value = param_cls(new_value, requires_grad=old_value.requires_grad, **kwargs).to(device)



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 244bf89b3..bd4b72682 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -127,13 +127,11 @@ def set_module_tensor_to_device(
 
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif module._parameters[tensor_name].device.type != "cuda":
+        elif torch.device(device) != module._parameters[tensor_name].device:
             param_cls = type(module._parameters[tensor_name])
             kwargs = module._parameters[tensor_name].__dict__
             new_value = param_cls(new_value, requires_grad=old_value.requires_grad, **kwargs).to(device)
             module._parameters[tensor_name] = new_value
-        else:
-            module._parameters[tensor_name] = module._parameters[tensor_name].to(device)
 
 
 def named_module_tensors(module: nn.Module, include_buffers: bool = True, recurse: bool = False):


Review comments: body: ```suggestion
        elif device != module._parameters[tensor_name].device.index:
```
The `device` param is an integer so to properly compare the devices we need to access the index of the `device` attribute of the parameter ;) 
diff_hunk: @@ -127,13 +127,11 @@ def set_module_tensor_to_device(
 
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif module._parameters[tensor_name].device.type != "cuda":
+        elif device != module._parameters[tensor_name].device:
body: Ah should be `torch.device(device)` on the left side to cover all cases!
diff_hunk: @@ -127,13 +127,11 @@ def set_module_tensor_to_device(
 
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif module._parameters[tensor_name].device.type != "cuda":
+        elif device != module._parameters[tensor_name].device:
body: Fixed [here](https://github.com/huggingface/accelerate/pull/576/commits/10334600357bdd014acd67bb255884ee10b54fe6) normally.
diff_hunk: @@ -127,13 +127,11 @@ def set_module_tensor_to_device(
 
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif module._parameters[tensor_name].device.type != "cuda":
+        elif device != module._parameters[tensor_name].device:
body: Can confirm this fixes the issue too !
diff_hunk: @@ -127,13 +127,11 @@ def set_module_tensor_to_device(
 
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif module._parameters[tensor_name].device.type != "cuda":
+        elif device != module._parameters[tensor_name].device:

Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/.gitignore b/.gitignore
index 7f145cd5c..da99824aa 100644
--- a/.gitignore
+++ b/.gitignore
@@ -135,4 +135,7 @@ dmypy.json
 .idea
 
 # Mac .DS_Store
-.DS_Store
\ No newline at end of file
+.DS_Store
+
+# More test things
+wandb
\ No newline at end of file
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 8983468e1..ac8c4de52 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -20,6 +20,8 @@
 import unittest
 from unittest import mock
 
+import torch
+
 from accelerate.test_utils.examples import compare_against_test
 from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
 from accelerate.utils import write_basic_config
@@ -169,8 +171,16 @@ def test_load_states_by_steps(self):
         --resume_from_checkpoint {os.path.join(self.tmpdir, "step_2")}
         """.split()
         output = run_command(self._launch_args + testargs, return_stdout=True)
-        self.assertNotIn("epoch 0:", output)
-        self.assertIn("epoch 1:", output)
+        if torch.cuda.is_available():
+            num_processes = torch.cuda.device_count()
+        else:
+            num_processes = 1
+        if num_processes > 1:
+            self.assertNotIn("epoch 0:", output)
+            self.assertIn("epoch 1:", output)
+        else:
+            self.assertIn("epoch 0:", output)
+            self.assertIn("epoch 1:", output)
 
     @slow
     def test_cross_validation(self):



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 361e19642..8983468e1 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -166,7 +166,7 @@ def test_load_states_by_epoch(self):
     def test_load_states_by_steps(self):
         testargs = f"""
         examples/by_feature/checkpointing.py
-        --resume_from_checkpoint {os.path.join(self.tmpdir, "step_5")}
+        --resume_from_checkpoint {os.path.join(self.tmpdir, "step_2")}
         """.split()
         output = run_command(self._launch_args + testargs, return_stdout=True)
         self.assertNotIn("epoch 0:", output)



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index a0f9fab40..b84a90c6e 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -29,10 +29,12 @@
 
 try:
     import torch_ccl  # noqa: F401
+
     _ccl_available = True
 except ImportError:
-    try: 
-        import oneccl_bindings_for_pytorch
+    try:
+        import oneccl_bindings_for_pytorch  # noqa: F401
+
         _ccl_available = True
     except ImportError:
         _ccl_available = False



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/.github/workflows/check_dependencies.yml b/.github/workflows/check_dependencies.yml
index 1e8a4e1bb..f8ac1f492 100644
--- a/.github/workflows/check_dependencies.yml
+++ b/.github/workflows/check_dependencies.yml
@@ -1,4 +1,4 @@
-name: Trigger docker images and run slow tests
+name: Trigger docker images and run tests
 
 on:
   push:
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index 0e8f31407..c0c9c62f2 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -1,4 +1,4 @@
-name: Self-hosted runner (scheduled)
+name: Self-hosted runner with slow tests (scheduled)
 
 on:
   workflow_dispatch:



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index c846dd65d..06a25bc4e 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -17,7 +17,7 @@
 
 import torch
 
-from .utils import DistributedType, is_ccl_available, is_deepspeed_available, is_tpu_available
+from .utils import DistributedType, get_ccl_version, is_ccl_available, is_deepspeed_available, is_tpu_available
 from .utils.dataclasses import SageMakerDistributedType
 
 
@@ -166,6 +166,10 @@ def __init__(
             elif get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1:
                 self.distributed_type = DistributedType.MULTI_CPU
                 if is_ccl_available() and get_int_from_env(["CCL_WORKER_COUNT"], 0) > 0:
+                    if get_ccl_version() >= "1.12":
+                        import oneccl_bindings_for_pytorch  # noqa: F401
+                    else:
+                        import torch_ccl  # noqa: F401
                     backend = "ccl"
                 elif torch.distributed.is_mpi_available():
                     backend = "mpi"
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index f6b83d2f9..0395a271d 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -19,6 +19,7 @@
     TensorInformation,
 )
 from .imports import (
+    get_ccl_version,
     is_apex_available,
     is_bf16_available,
     is_boto3_available,
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index b84a90c6e..da63d0aa1 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -27,19 +27,6 @@
     import importlib.metadata as importlib_metadata
 
 
-try:
-    import torch_ccl  # noqa: F401
-
-    _ccl_available = True
-except ImportError:
-    try:
-        import oneccl_bindings_for_pytorch  # noqa: F401
-
-        _ccl_available = True
-    except ImportError:
-        _ccl_available = False
-
-
 try:
     import torch_xla.core.xla_model as xm  # noqa: F401
 
@@ -49,7 +36,14 @@
 
 
 def is_ccl_available():
-    return _ccl_available
+    return (
+        importlib.util.find_spec("torch_ccl") is not None
+        or importlib.util.find_spec("oneccl_bindings_for_pytorch") is not None
+    )
+
+
+def get_ccl_version():
+    return importlib_metadata.version("oneccl_bind_pt")
 
 
 def is_apex_available():



Discussion: body: @liangan1 @yao-matrix @sgugger please have a review
body: _The documentation is not available anymore as the PR was closed or merged._
body: Thanks for this! Looks great to me 😄 could you run `make style; make quality` to get past that failure please? And then we're good!
body: @muellerzr thanks for the suggestion, fix the "make quality" issue.
body: done for the rebase @sgugger
diff --git a/setup.py b/setup.py
index 60b6647fd..cad067ad5 100644
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,7 @@
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3", "hf-doc-builder >= 0.3.0"]
 extras["docs"] = []
 extras["test_prod"] = ["pytest", "pytest-xdist", "pytest-subtests", "parameterized"]
-extras["test_dev"] = ["datasets<=2.2.2", "evaluate", "transformers", "scipy", "sklearn", "deepspeed"]
+extras["test_dev"] = ["datasets", "evaluate", "transformers", "scipy", "sklearn", "deepspeed"]
 extras["testing"] = extras["test_prod"] + extras["test_dev"]
 
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard"]



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 35afd6c5d..5aea696b4 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -446,16 +446,13 @@ def _prepare_one(self, obj, first_pass=False):
             if isinstance(obj, torch.utils.data.DataLoader):
                 return self.prepare_data_loader(obj)
             elif isinstance(obj, torch.nn.Module):
-                self._models.append(obj)
                 return self.prepare_model(obj)
             elif isinstance(obj, torch.optim.Optimizer):
                 optimizer = self.prepare_optimizer(obj)
-                self._optimizers.append(optimizer)
                 return optimizer
         # Second pass of preparation: LR scheduler (which need the full list of optimizers)
         elif isinstance(obj, torch.optim.lr_scheduler._LRScheduler):
             scheduler = self.prepare_scheduler(obj)
-            self._schedulers.append(scheduler)
             return scheduler
         # Return the unprocessed object if previous criteria was not met
         return obj
@@ -570,6 +567,7 @@ def prepare(self, *args):
         return result if len(result) > 1 else result[0]
 
     def prepare_model(self, model):
+        self._models.append(model)
         if self.device_placement and self.distributed_type != DistributedType.FSDP:
             model = model.to(self.device)
         if self.distributed_type == DistributedType.MULTI_GPU:
@@ -782,7 +780,9 @@ def prepare_data_loader(self, data_loader):
         )
 
     def prepare_optimizer(self, optimizer):
-        return AcceleratedOptimizer(optimizer, device_placement=self.device_placement, scaler=self.scaler)
+        optimizer = AcceleratedOptimizer(optimizer, device_placement=self.device_placement, scaler=self.scaler)
+        self._optimizers.append(optimizer)
+        return optimizer
 
     def prepare_scheduler(self, scheduler):
         # We try to find the optimizer associated with `scheduler`, the default is the full list.
@@ -791,13 +791,14 @@ def prepare_scheduler(self, scheduler):
             if getattr(scheduler, "optimizer", None) == opt.optimizer:
                 optimizer = opt
                 break
-
-        return AcceleratedScheduler(
+        scheduler = AcceleratedScheduler(
             scheduler,
             optimizer,
             step_with_optimizer=self.step_scheduler_with_optimizer,
             split_batches=self.split_batches,
         )
+        self._schedulers.append(scheduler)
+        return scheduler
 
     def backward(self, loss, **kwargs):
         """



Discussion: body: Thanks! If we want to go this route, it would be good to update the rest of them as well. I.e. optimizer and scheduler
body: _The documentation is not available anymore as the PR was closed or merged._
body: @muellerzr thanks for the comment, I updated the optimizer and scheduler as well. If any other changes are needed, please let me know.
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index a660db858..8bffd8e99 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -29,10 +29,13 @@
 
 try:
     import torch_ccl  # noqa: F401
-
     _ccl_available = True
 except ImportError:
-    _ccl_available = False
+    try: 
+        import oneccl_bindings_for_pytorch
+        _ccl_available = True
+     except ImportError:
+        _ccl_available = False
 
 
 try:
diff --git a/docs/source/accelerator.mdx b/docs/source/accelerator.mdx
index 5202aeb8c..268b09d1f 100644
--- a/docs/source/accelerator.mdx
+++ b/docs/source/accelerator.mdx
@@ -44,11 +44,12 @@ actually be performed, and auto-scale the loss:
 
 ```python
 accelerator = Accelerator(gradient_accumulation_steps=2)
+model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)
 
-for (input, label) in enumerate(training_dataloader):
+for input, label in training_dataloader:
     with accelerator.accumulate(model):
         predictions = model(input)
-        loss = loss_function(predictions, labels)
+        loss = loss_function(predictions, label)
         accelerator.backward(loss)
         optimizer.step()
         scheduler.step()



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
body: @sgugger Thanks, I applied the change.
body: Thanks a lot!
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 6d3e7d3ff..c846dd65d 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -130,8 +130,16 @@ def __init__(
                 ), "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
                 self.distributed_type = DistributedType.DEEPSPEED
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend="nccl", **kwargs)
+                    from .utils import compare_versions
+
                     self.backend = "nccl"
+                    if compare_versions("deepspeed", ">", "0.6.5"):
+                        from deepspeed import comm as dist
+
+                        dist.init_distributed(dist_backend=self.backend)
+                    else:
+                        torch.distributed.init_process_group(backend="nccl", **kwargs)
+
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
                 self.local_process_index = int(os.environ.get("LOCAL_RANK", -1))


Review comments: body: ```suggestion
    "deepspeed<=0.6.5",
```
I don't think we required 0.6.5 as a min version before, so let's not pin exactly.
diff_hunk: @@ -28,7 +28,7 @@
     "scipy",
     "sklearn",
     "parameterized",
-    "deepspeed",
+    "deepspeed==0.6.5",
body: This can be imported at the top.
diff_hunk: @@ -130,8 +130,16 @@ def __init__(
                 ), "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
                 self.distributed_type = DistributedType.DEEPSPEED
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend="nccl", **kwargs)
+                    from .utils import compare_versions
body: that was leading to circular dependency error

diff_hunk: @@ -130,8 +130,16 @@ def __init__(
                 ), "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
                 self.distributed_type = DistributedType.DEEPSPEED
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend="nccl", **kwargs)
+                    from .utils import compare_versions
body: Ah yes, I can see that! Good for me then!
diff_hunk: @@ -130,8 +130,16 @@ def __init__(
                 ), "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
                 self.distributed_type = DistributedType.DEEPSPEED
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend="nccl", **kwargs)
+                    from .utils import compare_versions

Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 5af38afaf..7f36121c1 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -434,22 +434,23 @@ def print(self, *args, **kwargs):
 
     def _prepare_one(self, obj, first_pass=False):
         # First pass of preparation: DataLoader, model, optimizer
-        if isinstance(obj, torch.utils.data.DataLoader) and first_pass:
-            return self.prepare_data_loader(obj)
-        elif isinstance(obj, torch.nn.Module) and first_pass:
-            self._models.append(obj)
-            return self.prepare_model(obj)
-        elif isinstance(obj, torch.optim.Optimizer) and first_pass:
-            optimizer = self.prepare_optimizer(obj)
-            self._optimizers.append(optimizer)
-            return optimizer
+        if first_pass:
+            if isinstance(obj, torch.utils.data.DataLoader):
+                return self.prepare_data_loader(obj)
+            elif isinstance(obj, torch.nn.Module):
+                self._models.append(obj)
+                return self.prepare_model(obj)
+            elif isinstance(obj, torch.optim.Optimizer):
+                optimizer = self.prepare_optimizer(obj)
+                self._optimizers.append(optimizer)
+                return optimizer
         # Second pass of preparation: LR scheduler (which need the full list of optimizers)
-        elif isinstance(obj, torch.optim.lr_scheduler._LRScheduler) and not first_pass:
+        elif isinstance(obj, torch.optim.lr_scheduler._LRScheduler):
             scheduler = self.prepare_scheduler(obj)
             self._schedulers.append(scheduler)
             return scheduler
-        else:
-            return obj
+        # Return the unprocessed object if previous criteria was not met
+        return obj
 
     def _prepare_fsdp(self, *args):
         result = []



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 85b1f3237..907dfdac4 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -260,7 +260,7 @@ def load_checkpoint_and_dispatch(
     offload_folder: Optional[Union[str, os.PathLike]] = None,
     offload_buffers: bool = False,
     dtype: Optional[Union[str, torch.dtype]] = None,
-    offload_state_dict: bool = False,
+    offload_state_dict: Optional[bool] = None,
     preload_module_classes: Optional[List[str]] = None,
 ):
     """
@@ -292,9 +292,10 @@ def load_checkpoint_and_dispatch(
             well as the parameters.
         dtype (`str` or `torch.dtype`, *optional*):
             If provided, the weights will be converted to that type when loaded.
-        offload_state_dict (`bool`, *optional*, defaults to `False`):
+        offload_state_dict (`bool`, *optional*):
             If `True`, will temporarily offload the CPU state dict on the hard drive to avoig getting out of CPU RAM if
-            the weight of the CPU state dict + the biggest shard does not fit.
+            the weight of the CPU state dict + the biggest shard does not fit. Will default to `True` if the device map
+            picked contains `"disk"` values.
         preload_module_classes (`List[str]`, *optional*):
             A list of classes whose instances should load all their weights (even in the submodules) at the beginning
             of the forward. This should only be used for classes that have submodules which are registered but not
@@ -305,6 +306,8 @@ def load_checkpoint_and_dispatch(
         device_map = infer_auto_device_map(
             model, max_memory=max_memory, no_split_module_classes=no_split_module_classes, dtype=dtype
         )
+    if offload_state_dict is None and "disk" in device_map.values():
+        offload_state_dict = True
     load_checkpoint_in_model(
         model,
         checkpoint,



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/commands/test.py b/src/accelerate/commands/test.py
index 379233679..41da75596 100644
--- a/src/accelerate/commands/test.py
+++ b/src/accelerate/commands/test.py
@@ -43,7 +43,7 @@ def test_command_parser(subparsers=None):
 
 
 def test_command(args):
-    script_name = os.path.sep.join(__file__.split(os.path.sep)[:-2] + ["test_utils", "test_script.py"])
+    script_name = os.path.sep.join(__file__.split(os.path.sep)[:-2] + ["test_utils", "scripts", "test_script.py"])
 
     test_args = f"""
         --config_file={args.config_file} {script_name}



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 2dec26a98..9b6e7f069 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -161,7 +161,10 @@ def __init__(
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
         self.logging_dir = logging_dir
-        self.log_with = filter_trackers(log_with, self.logging_dir)
+        trackers = filter_trackers(log_with, self.logging_dir)
+        if len(trackers) < 1 and log_with is not None:
+            warnings.warn(f"`log_with={log_with}` was passed but no supported trackers are currently installed.")
+        self.log_with = trackers
 
         if mixed_precision is not None:
             mixed_precision = str(mixed_precision)



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 4014e3972..85b1f3237 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -65,7 +65,9 @@ def init_empty_weights(include_buffers: bool = False):
     def register_empty_parameter(module, name, param):
         old_register_parameter(module, name, param)
         if param is not None:
-            module._parameters[name] = nn.Parameter(module._parameters[name].to(torch.device("meta")))
+            param_cls = type(module._parameters[name])
+            kwargs = module._parameters[name].__dict__
+            module._parameters[name] = param_cls(module._parameters[name].to(torch.device("meta")), **kwargs)
 
     def register_empty_buffer(module, name, buffer):
         old_register_buffer(module, name, buffer)


Review comments: body: Leaves those poor lines alone :-)
diff_hunk: @@ -39,9 +39,7 @@ def init_empty_weights(include_buffers: bool = False):
     Args:
         include_buffers (`bool`, *optional*, defaults to `False`):
             Whether or not to also put all buffers on the meta device while initializing.
-
     Example:
-
body: 🤣 yeah don't know what happened there
diff_hunk: @@ -39,9 +39,7 @@ def init_empty_weights(include_buffers: bool = False):
     Args:
         include_buffers (`bool`, *optional*, defaults to `False`):
             Whether or not to also put all buffers on the meta device while initializing.
-
     Example:
-
body: Fixed in aafea4a1e18ac6562bc678f3e607e0d2c5f89b38 ;) 
diff_hunk: @@ -39,9 +39,7 @@ def init_empty_weights(include_buffers: bool = False):
     Args:
         include_buffers (`bool`, *optional*, defaults to `False`):
             Whether or not to also put all buffers on the meta device while initializing.
-
     Example:
-

Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index f58752e53..a2c7a4a9b 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -463,7 +463,9 @@ def __iter__(self):
 
     def __len__(self):
         whole_length = super().__len__()
-        if self.drop_last:
+        if self.split_batches:
+            return whole_length
+        elif self.drop_last:
             return whole_length // self.state.num_processes
         else:
             return math.ceil(whole_length / self.state.num_processes)



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/setup.py b/setup.py
index 01dc0439a..a24428c33 100644
--- a/setup.py
+++ b/setup.py
@@ -22,7 +22,7 @@
     "pytest",
     "pytest-xdist",
     "pytest-subtests",
-    "datasets",
+    "datasets<=2.2.2",
     "evaluate",
     "transformers",
     "scipy",



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/setup.py b/setup.py
index b48dc63ad..01dc0439a 100644
--- a/setup.py
+++ b/setup.py
@@ -22,7 +22,7 @@
     "pytest",
     "pytest-xdist",
     "pytest-subtests",
-    "datasets<=2.3.1",
+    "datasets",
     "evaluate",
     "transformers",
     "scipy",



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/setup.py b/setup.py
index 01dc0439a..b48dc63ad 100644
--- a/setup.py
+++ b/setup.py
@@ -22,7 +22,7 @@
     "pytest",
     "pytest-xdist",
     "pytest-subtests",
-    "datasets",
+    "datasets<=2.3.1",
     "evaluate",
     "transformers",
     "scipy",



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
body: Looks like the PR with the fix will be merged soon though? But feel free to pin if you want in the meantime :-)
diff --git a/setup.py b/setup.py
index 94516436b..e548a3940 100644
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,7 @@
 from setuptools import find_packages
 
 extras = {}
-extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3"]
+extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3", "hf-doc-builder >= 0.3.0"]
 extras["docs"] = []
 extras["test"] = [
     "pytest",



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/docs/source/internal.mdx b/docs/source/internal.mdx
index 9b496b30e..4f7f59500 100644
--- a/docs/source/internal.mdx
+++ b/docs/source/internal.mdx
@@ -56,6 +56,12 @@ The main work on your PyTorch `DataLoader` is done by the following function:
 
 [[autodoc]] utils.extract_model_from_parallel
 
+[[autodoc]] utils.is_bf16_available
+
+[[autodoc]] utils.is_torch_version
+
+[[autodoc]] utils.is_tpu_available
+
 [[autodoc]] utils.gather
 
 [[autodoc]] utils.send_to_device
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index d1d634355..594f7aa1c 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -57,6 +57,7 @@ def is_apex_available():
 
 
 def is_tpu_available():
+    "Checks if `torch_xla` is installed and if a TPU is in the environment"
     return _tpu_available
 
 



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/.github/workflows/check_dependencies.yml b/.github/workflows/check_dependencies.yml
index c4818ba85..1e8a4e1bb 100644
--- a/.github/workflows/check_dependencies.yml
+++ b/.github/workflows/check_dependencies.yml
@@ -1,4 +1,4 @@
-name: Check for dependency modification
+name: Trigger docker images and run slow tests
 
 on:
   push:
@@ -37,4 +37,9 @@ jobs:
     needs: check-for-setup
     if: (github.event_name == 'push') && (needs.check-for-setup.outputs.changed == '1')
     uses: ./.github/workflows/build-docker-images.yml
-    secrets: inherit
\ No newline at end of file
+    secrets: inherit
+
+  run-tests:
+    needs: build-docker-containers
+    if: always()
+    uses: ./.github/workflows/on-merge.yml
\ No newline at end of file
diff --git a/.github/workflows/on-merge.yml b/.github/workflows/on-merge.yml
index f8b9ccddb..b95474758 100644
--- a/.github/workflows/on-merge.yml
+++ b/.github/workflows/on-merge.yml
@@ -1,10 +1,7 @@
-name: Self-hosted runner (push to "main")
+name: Self-hosted runner tests (push to "main")
 
 on:
-  workflow_run:
-    workflows: ["Check for dependency modification"]
-    branches: ["main"]
-    types: [completed]
+  workflow_call:
   workflow_dispatch:
 
 env:



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index c4487f6b6..2d2a0bd56 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -254,7 +254,7 @@ def __init__(
             self.scaler = torch.cuda.amp.GradScaler(**kwargs)
         elif self.state.mixed_precision == "bf16":
             self.native_amp = is_bf16_available(True)
-            if mixed_precision == "bf16" and not self.native_amp:
+            if mixed_precision == "bf16" and not self.native_amp and not is_tpu_available():
                 raise ValueError(err.format(mode="bf16", requirement="PyTorch >= 1.10 and a supported device."))
 
             # Only on the GPU do we care about scaling the gradients
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 95b9eb30a..36b11be56 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -80,7 +80,9 @@ def __init__(
         if not getattr(self, "initialized", False):
             self.backend = None
             self.deepspeed_plugin = None
-            mixed_precision = mixed_precision.lower() if mixed_precision else None
+            mixed_precision = (
+                parse_choice_from_env("MIXED_PRECISION", "no") if mixed_precision is None else mixed_precision.lower()
+            )
             if not _from_accelerator:
                 raise ValueError(
                     "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` "



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/utils/offload.py b/src/accelerate/utils/offload.py
index 6e8b34bbc..750ff9d57 100644
--- a/src/accelerate/utils/offload.py
+++ b/src/accelerate/utils/offload.py
@@ -25,8 +25,8 @@ def offload_weight(weight, weight_name, offload_folder, index=None):
     dtype = None
     # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
     if str(weight.dtype) == "torch.bfloat16":
-        # Need to convert to FP32 since NumPy does not handle bfloat16s.
-        weight = weight.float()
+        # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
+        weight = weight.view(torch.int16)
         dtype = "bfloat16"
     array = weight.numpy()
     tensor_file = os.path.join(offload_folder, f"{weight_name}.dat")
@@ -50,8 +50,8 @@ def load_offloaded_weight(weight_file, weight_info):
 
     dtype = weight_info["dtype"]
     if dtype == "bfloat16":
-        # NumPy does not support bfloat16 so this was saved as a float32
-        dtype = "float32"
+        # NumPy does not support bfloat16 so this was saved as a int16
+        dtype = "int16"
 
     weight = np.memmap(weight_file, dtype=dtype, shape=shape, mode="r")
 
@@ -59,7 +59,7 @@ def load_offloaded_weight(weight_file, weight_info):
         weight = weight[0]
     weight = torch.tensor(weight)
     if weight_info["dtype"] == "bfloat16":
-        weight = weight.to(torch.bfloat16)
+        weight = weight.view(torch.bfloat16)
 
     return weight
 



Discussion: body: This PR handles a bug reported in https://github.com/huggingface/accelerate/issues/454 concerning bfloat16 and disk offload.
body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index bb8aba326..38da07aa1 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -157,7 +157,7 @@ def training_function(config, args):
             # We use the new `no_sync` context manager to prevent gradient averaging
             # until we want to at the proper step if we happen to be in a distributed setup
             # otherwise it does nothing
-            if step % gradient_accumulation_steps == 0:
+            if step % gradient_accumulation_steps != 0:
                 # Accumulate gradients locally
                 with accelerator.no_sync(model):
                     output = model(**batch)



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index c8cce9587..0a44a093f 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -34,7 +34,12 @@
 try:
     import torch_xla.core.xla_model as xm  # noqa: F401
 
-    _tpu_available = True
+    try:
+        # Will raise a RuntimeError if no XLA configuration is found
+        _ = xm.xla_device()
+        _tpu_available = True
+    except RuntimeError:
+        _tpu_available = False
 except ImportError:
     _tpu_available = False
 



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index b043d5d20..8b41c48f8 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -116,7 +116,7 @@ def load_accelerator_state(input_dir, models, optimizers, schedulers, process_in
     Args:
         input_dir (`str` or `os.PathLike`):
             The name of the folder to load all relevant weights and states.
-        model_stmodelsates (`List[torch.nn.Module]`):
+        models (`List[torch.nn.Module]`):
             A list of model instances
         optimizers (`List[torch.optim.Optimizer]`):
             A list of optimizer instances



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/setup.py b/setup.py
index 45c4ae912..93916bd9f 100644
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,6 @@
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3"]
 extras["docs"] = []
 extras["test"] = [
-    "psutil",
     "pytest",
     "pytest-xdist",
     "pytest-subtests",
@@ -60,7 +59,7 @@
         ]
     },
     python_requires=">=3.7.0",
-    install_requires=["numpy>=1.17", "packaging>=20.0", "pyyaml", "torch>=1.4.0"],
+    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0"],
     extras_require=extras,
     classifiers=[
         "Development Status :: 5 - Production/Stable",



Discussion: body: Thanks! LG2M!
body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 8801011ec..407b90184 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -199,6 +199,10 @@ def __init__(
                     ds_config.update({"bf16": {"enabled": True}})
                 self.dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive # noqa
 
+        if os.environ.get("USE_FSDP", "false") == "true" or isinstance(fsdp_plugin, FullyShardedDataParallelPlugin):
+            if is_torch_version("<", "1.12.0.dev20220418+cu113"):
+                raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
+
         if fsdp_plugin is None:  # init from env variables
             fsdp_plugin = FullyShardedDataParallelPlugin() if os.environ.get("USE_FSDP", "false") == "true" else None
         else:
@@ -206,10 +210,6 @@ def __init__(
                 raise TypeError("`fsdp_plugin` must be a FullyShardedDataParallelPlugin object.")
             os.environ["USE_FSDP"] = "true"  # use FSDP if plugin is provided
 
-        if os.environ.get("USE_FSDP", "false") == "true":
-            if is_torch_version("<", "1.12.0.dev20220418+cu113"):
-                raise ValueError("FSDP requires PyTorch >= 1.12.0.dev20220418+cu113")
-
         # Kwargs handlers
         self.ddp_handler = None
         self.scaler_handler = None



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
index dbe64b5c7..9ece7b0f7 100644
--- a/tests/deepspeed/test_deepspeed.py
+++ b/tests/deepspeed/test_deepspeed.py
@@ -158,7 +158,6 @@ def test_deepspeed_plugin(self, stage):
             self.assertFalse(deepspeed_plugin.zero3_init_flag)
         elif stage == ZERO3:
             self.assertTrue(deepspeed_plugin.zero3_init_flag)
-            deepspeed_plugin.deepspeed_config = None
 
         # Test `gradient_accumulation_steps` is set to 1 if unavailable in config file
         with tempfile.TemporaryDirectory() as dirpath:
@@ -232,6 +231,10 @@ def test_deepspeed_plugin(self, stage):
 
     @parameterized.expand([FP16, BF16], name_func=parameterized_custom_name_func)
     def test_accelerate_state_deepspeed(self, dtype):
+        state = AcceleratorState(_from_accelerator=True)
+        if state.initialized:
+            state.initialized = False
+
         deepspeed_plugin = DeepSpeedPlugin(
             gradient_accumulation_steps=1,
             gradient_clipping=1.0,



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/.github/workflows/check_dependencies.yml b/.github/workflows/check_dependencies.yml
index 65afdd6c9..c4818ba85 100644
--- a/.github/workflows/check_dependencies.yml
+++ b/.github/workflows/check_dependencies.yml
@@ -36,4 +36,5 @@ jobs:
   build-docker-containers:
     needs: check-for-setup
     if: (github.event_name == 'push') && (needs.check-for-setup.outputs.changed == '1')
-    uses: ./.github/workflows/build-docker-images.yml
\ No newline at end of file
+    uses: ./.github/workflows/build-docker-images.yml
+    secrets: inherit
\ No newline at end of file



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
body: @muellerzr Thank you pining for me. I read (part of) that page previously, but not in detail.

So it looks like to me that the called (reusable) workflow can't access `secrets`, and need to get it from the caller.
(Only the workflows triggered by events like `push/pull_request/schedule`, etc.  can access `secrets`)
body: Yup exactly. And we can do so by using `secrets`, by specifying inherit it can have all of the secrets the caller has. Keeping it to be what's needed only isn't a bad idea either, rather than giving them everything. But for this one I decided all wouldn't be a terrible idea.
diff --git a/setup.py b/setup.py
index 45c4ae912..eb0f50ab8 100644
--- a/setup.py
+++ b/setup.py
@@ -14,7 +14,7 @@
 
 from setuptools import setup
 from setuptools import find_packages
-
+# A comment
 extras = {}
 extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3"]
 extras["docs"] = []



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index 1764e2381..d52cf5e1e 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -5,12 +5,14 @@ on:
   schedule:
     - cron: "0 2 * * *"
 
+env:
+  RUN_SLOW: "yes"
+
 jobs:
   run_all_tests_single_gpu:
     runs-on: [self-hosted, docker-gpu, multi-gpu]
     env:
       CUDA_VISIBLE_DEVICES: "0"
-      ALL_SLOW: "yes"
     container:
       image: huggingface/accelerate-gpu:latest
       options: --gpus all --shm-size "16gb"
@@ -38,7 +40,6 @@ jobs:
   run_all_tests_multi_gpu:
     runs-on: [self-hosted, docker-gpu, multi-gpu]
     env:
-      ALL_SLOW: "yes"
       CUDA_VISIBLE_DEVICES: "0,1"
     container:
       image: huggingface/accelerate-gpu:latest



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index d764d3927..2bd407d85 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -2,33 +2,41 @@ name: Run Tests
 
 on: [pull_request]
 
+env:
+  HF_HOME: ~/hf_cache
+
 jobs:
-  test:
+  run-tests:
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        test-kind: [
+          test,
+          test_example_differences,
+          test_checkpoint_step,
+          test_checkpoint_epoch,
+          test_rest
+        ]
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.7
-      uses: actions/setup-python@v2
+    - uses: actions/checkout@v3
+    - name: Set up python 3.7
+      uses: actions/setup-python@v3
       with:
         python-version: 3.7
-    - name: Install Python dependencies
-      run: |
-        pip install --upgrade pip
-        pip install -e .[test,test_trackers]
-    - name: Run Tests
-      run: make test
-      
-  test_examples:
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.7
-      uses: actions/setup-python@v2
+    
+    - name: Activate python cache
+      uses: actions/cache@v3
       with:
-        python-version: 3.7
-    - name: Install Python dependencies
+        path: |
+          ${{ env.pythonLocation }}
+          ${{ env.HF_HOME }}
+        key: ${{ env.pythonLocation }}-${{ matrix.test-kind }}-${{ hashFiles('setup.py') }}
+    
+    - name: Install the library
       run: |
         pip install --upgrade pip
-        pip install -e .[test,tensorboard]
+        pip install -e .[test,test_trackers,tensorboard]
+    
     - name: Run Tests
-      run: make test_examples
+      run: |
+        make ${{ matrix.test-kind }}
\ No newline at end of file
diff --git a/Makefile b/Makefile
index f4534e12c..38cb16cfc 100644
--- a/Makefile
+++ b/Makefile
@@ -29,3 +29,16 @@ test:
 
 test_examples:
 	python -m pytest -s -v ./tests/test_examples.py
+
+# Broken down example tests for the CI runners
+test_example_differences:
+	python -m pytest -s -v ./tests/test_examples.py::ExampleDifferenceTests
+
+test_checkpoint_epoch:
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_epoch"
+
+test_checkpoint_step:
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_step"
+
+test_rest:
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "not by_step and not by_epoch"
\ No newline at end of file



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index 11946b544..d50aab329 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -32,7 +32,7 @@ jobs:
       - name: Run test on GPUs
         run: |
           source activate accelerate
-          make test_cuda
+          make test
       - name: Run examples on GPUs
         run: |
           source activate accelerate
@@ -64,7 +64,7 @@ jobs:
       - name: Run test on GPUs
         run: |
           source activate accelerate
-          make test_cuda
+          make test
 
       - name: Run examples on GPUs
         run: |
diff --git a/.github/workflows/on-merge.yml b/.github/workflows/on-merge.yml
new file mode 100644
index 000000000..8df455563
--- /dev/null
+++ b/.github/workflows/on-merge.yml
@@ -0,0 +1,69 @@
+name: Self-hosted runner (push to "main")
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - "main"
+
+jobs:
+  run_all_tests_single_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install tensorflow -U 
+          pip install -e . -U
+
+# Note: tensorflow upgrade is needed until dropped py 3.6 support
+
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          make test_examples
+
+  run_all_tests_multi_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install tensorflow -U
+          pip install -e . -U
+
+        # Note: tensorflow upgrade is needed until dropped py 3.6 support
+
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test
+
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          make test_examples
\ No newline at end of file



Discussion: body: Thinking on one of your frustrations @sgugger, I decided to add a `pip install -U` instead of `--no-deps` so that if a dep was added on a push, we can test it (and not instead have all those failing tests on a push to main and freak out 😉 )
body: _The documentation is not available anymore as the PR was closed or merged._