set os environ RANK

src/axolotl/utils/models.py

@@ -828,6 +828,7 @@ class ModelLoader:
 
             if self.cfg.tensor_parallel == "auto":
                 rank = int(os.environ.get("LOCAL_RANK", 0))
+                os.environ["RANK"] = str(rank)
                 device = torch.device(f"cuda:{rank}")
                 torch.distributed.init_process_group("nccl", device_id=device)