added logging to examples

Author: apoorvkh

scripts/examples/accelerate_train.py

@@ -13,6 +13,7 @@
 from __future__ import annotations
 
 import functools
+import logging
 import os
 from dataclasses import dataclass
 from pathlib import Path
@@ -27,6 +28,8 @@
 
 import torchrunx
 
+logging.basicConfig(level=logging.INFO)
+
 
 @dataclass
 class ModelConfig:
@@ -114,14 +117,18 @@ def main(
     output_dir: Path,
 ):
     model = AutoModelForCausalLM.from_pretrained(model_config.name)
-    train_dataset = load_training_data(tokenizer_name=model_config.name, dataset_config=dataset_config)
+    train_dataset = load_training_data(
+        tokenizer_name=model_config.name, dataset_config=dataset_config
+    )
 
     # Launch training
     results = launcher.run(train, model, train_dataset, batch_size, output_dir)
 
     # Loading trained model from checkpoint
     checkpoint_path = results.rank(0)
-    trained_model = AutoModelForCausalLM.from_pretrained(model_config.name, state_dict=torch.load(checkpoint_path))
+    trained_model = AutoModelForCausalLM.from_pretrained(
+        model_config.name, state_dict=torch.load(checkpoint_path)
+    )
 
 
 if __name__ == "__main__":

scripts/examples/deepspeed_train.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import functools
+import logging
 import os
 from dataclasses import dataclass
 from pathlib import Path
@@ -30,6 +31,8 @@
 
 import torchrunx
 
+logging.basicConfig(level=logging.INFO)
+
 
 @dataclass
 class DatasetConfig:

scripts/examples/lightning_train.py

@@ -14,13 +14,13 @@
 from __future__ import annotations
 
 import functools
+import logging
 import os
 from dataclasses import dataclass
 from typing import Annotated
 
 import lightning as L
 import torch
-
 import tyro
 from datasets import load_dataset
 from torch.utils.data import Dataset
@@ -29,6 +29,8 @@
 import torchrunx
 from torchrunx.integrations.lightning import TorchrunxClusterEnvironment
 
+logging.basicConfig(level=logging.INFO)
+
 
 @dataclass
 class ModelConfig:

scripts/examples/transformers_train.py

@@ -13,6 +13,7 @@
 from __future__ import annotations
 
 import functools
+import logging
 import os
 from dataclasses import dataclass
 from typing import Annotated
@@ -30,6 +31,7 @@
 
 import torchrunx
 
+logging.basicConfig(level=logging.INFO)
 
 @dataclass
 class ModelConfig: