Merge branch 'update-docs' of https://github.com/apoorvkh/torchrunx into update-docs

Author: apoorvkh

docs/source/examples/lightning_example.py

@@ -10,7 +10,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 import torchrunx
-
+from torchrunx.ext.lightning import TorchrunxClusterEnvironment
 
 class GPT2CausalLMDataset(Dataset):
     def __init__(self, text_dataset):
@@ -47,7 +47,7 @@ def __init__(self):
         super().__init__()
         self.model = AutoModelForCausalLM.from_pretrained("gpt2")
 
-    def training_step(self, batch, batch_idx):
+    def training_step(self, batch, *args): # pyright: ignore
         device_batch = {k: v.to(self.model.device) for k, v in batch.items()}
         loss = self.model(**device_batch).loss
         self.log("train_loss", loss)
@@ -72,24 +72,23 @@ def train():
         devices=2,
         num_nodes=1,
         strategy="ddp",
+        plugins=[TorchrunxClusterEnvironment()],
+        enable_checkpointing=False
     )
 
     trainer.fit(model=lightning_model, train_dataloaders=train_loader)
+    checkpoint  = f"{trainer.log_dir}/final.ckpt"
+    trainer.save_checkpoint(checkpoint)
 
-    if int(os.environ["RANK"]) == 0:
-        return trainer.model.model
-    return None
+    return checkpoint
 
 
 if __name__ == "__main__":
-    # hack to prevent lightning from recognizing SLURM environment...
-    os.environ["SLURM_JOB_NAME"] = "bash"
-    Path("output").mkdir(exist_ok=True)
     results = torchrunx.launch(
         func=train,
         hostnames=["localhost"],
         workers_per_host=2,
     )
 
-    trained_model: nn.Module = results.rank(0)
-    torch.save(trained_model.state_dict(), "output/model.pth")
+    checkpoint_path = results.rank(0)
+    print(f"Checkpoint at: {checkpoint_path}")

src/torchrunx/agent.py

@@ -83,6 +83,7 @@ def main(launcher_agent_group: LauncherAgentGroup, logger_hostname: str, logger_
                     backend=launcher_payload.backend,
                     rank=worker_global_ranks[i],
                     local_rank=i,
+                    node_rank=agent_rank,
                     local_world_size=num_workers,
                     world_size=worker_world_size,
                     hostname=launcher_payload.hostnames[agent_rank],

src/torchrunx/ext/__init__.py

@@ -0,0 +1 @@
+"""Extensions classes and functions."""

src/torchrunx/ext/lightning.py

@@ -0,0 +1,15 @@
+"""Pytorch Lightning extension utilities."""
+
+import torch
+from lightning.fabric.plugins.environments.torchelastic import (  # pyright: ignore [reportMissingImports]
+    TorchElasticEnvironment,
+)
+
+
+class TorchrunxClusterEnvironment(TorchElasticEnvironment):
+    """PyTorch Lightning ClusterEnvironment compatible with torchrunx."""
+
+    @staticmethod
+    def detect() -> bool:
+        """Returns ``True`` if the current process was launched using torchrunx."""
+        return torch.distributed.is_available()

src/torchrunx/worker.py

@@ -32,6 +32,7 @@ class WorkerArgs:
     backend: Literal["nccl", "gloo", "mpi", "ucc", "auto"] | None
     rank: int
     local_rank: int
+    node_rank: int
     local_world_size: int
     world_size: int
     hostname: str
@@ -79,6 +80,7 @@ def worker_entrypoint(serialized_worker_args: SerializedWorkerArgs) -> Any | Exc
 
     os.environ["RANK"] = str(worker_args.rank)
     os.environ["LOCAL_RANK"] = str(worker_args.local_rank)
+    os.environ["GROUP_RANK"] = str(worker_args.node_rank)
     os.environ["LOCAL_WORLD_SIZE"] = str(worker_args.local_world_size)
     os.environ["WORLD_SIZE"] = str(worker_args.world_size)
     os.environ["MASTER_ADDR"] = worker_args.main_agent_hostname