Update requirements for Ligthning Deepspeed (#9)

Author: ankitgola005

Lightning/DeepSpeed_Lightning/DeepSpeed_Lightning.ipynb

@@ -313,7 +313,7 @@
     }
    ],
    "source": [
-    "!pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.11.0"
+    "!pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.14.0"
    ]
   },
   {

Lightning/DeepSpeed_Lightning/lightning_gpt/models.py

@@ -199,7 +199,7 @@ def training_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor:
         return loss
 
     def generate(
-        self, idx: torch.Tensor, max_new_tokens: int, temperature: float = 1.0, top_k: Optional[int] = None
+        self, idx: torch.Tensor, max_new_tokens: int, temperature: float = 1.0, do_sample:bool = False, top_k: Optional[int] = None
     ) -> torch.Tensor:
         return self.nanogpt.generate(idx, max_new_tokens, temperature, top_k)
 
@@ -247,7 +247,7 @@ def configure_optimizers(self) -> torch.optim.AdamW:
 
 class DeepSpeedNanoGPT(NanoGPT):
     # TODO: activation checkpointing (requires overriding forward)
-    def __init__(self, fused_adam: bool = True, offload: bool = False, **kwargs: Any):
+    def __init__(self, fused_adam: bool = False, offload: bool = False, **kwargs: Any):
         if fused_adam and offload:
             raise RuntimeError(
                 "Cannot use FusedAdam and CPUAdam at the same time! "

Lightning/DeepSpeed_Lightning/requirements.txt

@@ -1,4 +1,4 @@
-lightning>=1.8.0
-torch>=1.10.0
+lightning>=2.2.0
+torch>=2.2.0
 numpy
-lightning-habana==1.1.0
+lightning-habana==1.4.0

Lightning/DeepSpeed_Lightning/train.py

@@ -7,10 +7,10 @@
 
 if module_available("lightning"):
     import lightning.pytorch as L
-    from lightning.pytorch.plugins import DeepSpeedPrecisionPlugin
+    from lightning.pytorch.plugins import DeepSpeedPrecision
 elif module_available("pytorch_lightning"):
     import pytorch_lightning as L
-    from pytorch_lightning.plugins import DeepSpeedPrecisionPlugin
+    from pytorch_lightning.plugins import DeepSpeedPrecision
 
 import gc
 
@@ -157,7 +157,6 @@ def main(args):
         else HPUParallelStrategy(bucket_cap_mb=125, gradient_as_bucket_view=True, static_graph=True),
         callbacks=callback_list,
         accumulate_grad_batches=1,
-        precision="bf16-mixed" if args.strategy == "deepspeed" else "16-mixed",  # 16,
         max_epochs=args.max_epochs,
         num_nodes=1,
         check_val_every_n_epoch=5000,
@@ -166,7 +165,7 @@ def main(args):
         limit_val_batches=10,
         max_steps=args.max_steps,
         gradient_clip_val=1.0,
-        plugins=[DeepSpeedPrecisionPlugin(precision="bf16-mixed")] if args.strategy == "deepspeed" else None,
+        plugins=[DeepSpeedPrecision(precision="bf16-mixed")] if args.strategy == "deepspeed" else None,
     )
 
     trainer.fit(model, train_loader)