Support device_map=sequential & max_memory config parameters (#903)

* Support device_map sequential (and others). Support max_memory in cfg.

* Update documentation in README accordingly.

* Update README.md

---------

Co-authored-by: Wing Lian <wing.lian@gmail.com>

src/axolotl/utils/config.py

@@ -27,7 +27,7 @@ def choose_device(cfg):
 
     cfg.device = get_device()
     if cfg.world_size == 1:
-        cfg.device_map = "auto"
+        cfg.device_map = cfg.device_map or "auto"
     else:
         if cfg.device.startswith("cuda"):
             cfg.device_map = {"": torch.cuda.current_device()}

src/axolotl/utils/models.py

@@ -216,6 +216,7 @@ def load_model(
     model_kwargs = {}
 
     model_kwargs["device_map"] = cfg.device_map
+    model_kwargs["max_memory"] = cfg.max_memory
     model_kwargs["torch_dtype"] = cfg.torch_dtype
 
     if cfg.model_revision: