adding logger args

extensions/thunder/pretrain.py

@@ -5,6 +5,7 @@
 import pprint
 import sys
 import time
+from dataclasses import asdict
 from datetime import timedelta
 from functools import partial
 from pathlib import Path
@@ -20,7 +21,7 @@
 from typing_extensions import Literal
 
 from litgpt import Tokenizer
-from litgpt.args import EvalArgs, TrainArgs
+from litgpt.args import EvalArgs, LogArgs, TrainArgs
 from litgpt.data import DataModule, TinyLlama
 from litgpt.model import GPT, Block, CausalSelfAttention, Config, LLaMAMLP
 from litgpt.utils import (
@@ -70,6 +71,7 @@ def setup(
         tie_embeddings=False,
     ),
     eval: EvalArgs = EvalArgs(interval=1000, max_iters=100),
+    log: LogArgs = LogArgs(),
     optimizer: Union[str, Dict] = "AdamW",
     devices: Union[int, str] = "auto",
     num_nodes: int = 1,
@@ -121,7 +123,12 @@ def setup(
     tokenizer = Tokenizer(tokenizer_dir) if tokenizer_dir is not None else None
 
     logger = choose_logger(
-        logger_name, out_dir, name=f"pretrain-{config.name}", resume=bool(resume), log_interval=train.log_interval
+        logger_name,
+        out_dir,
+        name=f"pretrain-{config.name}",
+        resume=bool(resume),
+        log_interval=train.log_interval,
+        log_args=asdict(log),
     )
 
     if devices * num_nodes > 1:

litgpt/args.py

@@ -90,3 +90,15 @@ class EvalArgs:
     evaluate_example: Union[str, int] = "first"
     """How to pick an example instruction to evaluate periodically during training.
        Can be "first", "random", or an integer index to pick a specific example."""
+
+
+@dataclass
+class LogArgs:
+    """Logging-related arguments"""
+
+    project: Optional[str] = None
+    """Project name"""
+    run: Optional[str] = None
+    """Run name"""
+    group: Optional[str] = None
+    """Group name"""

litgpt/finetune/adapter.py

@@ -18,7 +18,7 @@
 from torchmetrics import RunningMean
 
 from litgpt.adapter import GPT, Block, Config, adapter_filter, mark_only_adapter_as_trainable
-from litgpt.args import EvalArgs, TrainArgs
+from litgpt.args import EvalArgs, LogArgs, TrainArgs
 from litgpt.data import Alpaca, DataModule
 from litgpt.generate.base import generate
 from litgpt.prompts import save_prompt_style
@@ -62,6 +62,7 @@ def setup(
         max_seq_length=None,
     ),
     eval: EvalArgs = EvalArgs(interval=100, max_new_tokens=100, max_iters=100),
+    log: LogArgs = LogArgs(),
     optimizer: Union[str, Dict] = "AdamW",
     logger_name: Literal["wandb", "tensorboard", "csv", "mlflow"] = "csv",
     seed: int = 1337,
@@ -95,7 +96,13 @@ def setup(
     config = Config.from_file(checkpoint_dir / "model_config.yaml")
 
     precision = precision or get_default_supported_precision(training=True)
-    logger = choose_logger(logger_name, out_dir, name=f"finetune-{config.name}", log_interval=train.log_interval)
+    logger = choose_logger(
+        logger_name,
+        out_dir,
+        name=f"finetune-{config.name}",
+        log_interval=train.log_interval,
+        log_args=dataclasses.asdict(log),
+    )
 
     plugins = None
     if quantize is not None and quantize.startswith("bnb."):

litgpt/finetune/adapter_v2.py

@@ -18,7 +18,7 @@
 from torchmetrics import RunningMean
 
 from litgpt.adapter_v2 import GPT, Block, Config, adapter_filter, mark_only_adapter_v2_as_trainable
-from litgpt.args import EvalArgs, TrainArgs
+from litgpt.args import EvalArgs, LogArgs, TrainArgs
 from litgpt.data import Alpaca, DataModule
 from litgpt.generate.base import generate
 from litgpt.prompts import save_prompt_style
@@ -64,6 +64,7 @@ def setup(
         max_seq_length=None,
     ),
     eval: EvalArgs = EvalArgs(interval=100, max_new_tokens=100, max_iters=100),
+    log: LogArgs = LogArgs(),
     optimizer: Union[str, Dict] = "AdamW",
     logger_name: Literal["wandb", "tensorboard", "csv", "mlflow"] = "csv",
     seed: int = 1337,
@@ -97,7 +98,13 @@ def setup(
     config = Config.from_file(checkpoint_dir / "model_config.yaml")
 
     precision = precision or get_default_supported_precision(training=True)
-    logger = choose_logger(logger_name, out_dir, name=f"finetune-{config.name}", log_interval=train.log_interval)
+    logger = choose_logger(
+        logger_name,
+        out_dir,
+        name=f"finetune-{config.name}",
+        log_interval=train.log_interval,
+        log_args=dataclasses.asdict(log),
+    )
 
     plugins = None
     if quantize is not None and quantize.startswith("bnb."):

litgpt/finetune/full.py

@@ -13,7 +13,7 @@
 from torch.utils.data import ConcatDataset, DataLoader
 from torchmetrics import RunningMean
 
-from litgpt.args import EvalArgs, TrainArgs
+from litgpt.args import EvalArgs, LogArgs, TrainArgs
 from litgpt.data import Alpaca, DataModule
 from litgpt.generate.base import generate
 from litgpt.model import GPT, Block, Config
@@ -58,6 +58,7 @@ def setup(
         max_seq_length=None,
     ),
     eval: EvalArgs = EvalArgs(interval=600, max_new_tokens=100, max_iters=100),
+    log: LogArgs = LogArgs(),
     optimizer: Union[str, Dict] = "AdamW",
     logger_name: Literal["wandb", "tensorboard", "csv", "mlflow"] = "csv",
     seed: int = 1337,
@@ -94,7 +95,12 @@ def setup(
 
     precision = precision or get_default_supported_precision(training=True)
     logger = choose_logger(
-        logger_name, out_dir, name=f"finetune-{config.name}", resume=bool(resume), log_interval=train.log_interval
+        logger_name,
+        out_dir,
+        name=f"finetune-{config.name}",
+        resume=bool(resume),
+        log_interval=train.log_interval,
+        log_args=dataclasses.asdict(log),
     )
 
     if devices * num_nodes > 1:

litgpt/finetune/lora.py

@@ -17,7 +17,7 @@
 from torch.utils.data import ConcatDataset, DataLoader
 from torchmetrics import RunningMean
 
-from litgpt.args import EvalArgs, TrainArgs
+from litgpt.args import EvalArgs, LogArgs, TrainArgs
 from litgpt.data import Alpaca, DataModule
 from litgpt.generate.base import generate
 from litgpt.lora import GPT, Block, Config, lora_filter, mark_only_lora_as_trainable
@@ -71,6 +71,7 @@ def setup(
         epochs=5,
         max_seq_length=None,
     ),
+    log: LogArgs = LogArgs(),
     eval: EvalArgs = EvalArgs(interval=100, max_new_tokens=100, max_iters=100),
     optimizer: Union[str, Dict] = "AdamW",
     logger_name: Literal["wandb", "tensorboard", "csv", "mlflow"] = "csv",
@@ -125,7 +126,13 @@ def setup(
     )
 
     precision = precision or get_default_supported_precision(training=True)
-    logger = choose_logger(logger_name, out_dir, name=f"finetune-{config.name}", log_interval=train.log_interval)
+    logger = choose_logger(
+        logger_name,
+        out_dir,
+        name=f"finetune-{config.name}",
+        log_interval=train.log_interval,
+        log_args=dataclasses.asdict(log),
+    )
 
     plugins = None
     if quantize is not None and quantize.startswith("bnb."):

litgpt/pretrain.py

@@ -3,6 +3,7 @@
 import math
 import pprint
 import time
+from dataclasses import asdict
 from datetime import timedelta
 from functools import partial
 from pathlib import Path
@@ -18,7 +19,7 @@
 from typing_extensions import Literal
 
 from litgpt import Tokenizer
-from litgpt.args import EvalArgs, TrainArgs
+from litgpt.args import EvalArgs, LogArgs, TrainArgs
 from litgpt.config import name_to_config
 from litgpt.data import DataModule, TinyLlama
 from litgpt.model import GPT, Block, CausalSelfAttention, Config, LLaMAMLP
@@ -62,6 +63,7 @@ def setup(
         tie_embeddings=False,
     ),
     eval: EvalArgs = EvalArgs(interval=1000, max_iters=100),
+    log: LogArgs = LogArgs(),
     optimizer: Union[str, Dict] = "AdamW",
     devices: Union[int, str] = "auto",
     num_nodes: int = 1,
@@ -127,7 +129,12 @@ def setup(
     tokenizer = Tokenizer(tokenizer_dir) if tokenizer_dir is not None else None
 
     logger = choose_logger(
-        logger_name, out_dir, name=f"pretrain-{config.name}", resume=bool(resume), log_interval=train.log_interval
+        logger_name,
+        out_dir,
+        name=f"pretrain-{config.name}",
+        resume=bool(resume),
+        log_interval=train.log_interval,
+        log_args=asdict(log),
     )
 
     if devices * num_nodes > 1:

litgpt/utils.py

@@ -542,6 +542,7 @@ def choose_logger(
     out_dir: Path,
     name: str,
     log_interval: int = 1,
+    log_args: Optional[Dict] = None,
     resume: Optional[bool] = None,
     **kwargs: Any,
 ):
@@ -550,7 +551,10 @@ def choose_logger(
     if logger_name == "tensorboard":
         return TensorBoardLogger(root_dir=(out_dir / "logs"), name="tensorboard", **kwargs)
     if logger_name == "wandb":
-        return WandbLogger(project=name, resume=resume, **kwargs)
+        project = log_args.pop("project", name)
+        run = log_args.pop("run", os.environ.get("WANDB_RUN_NAME"))
+        group = log_args.pop("group", os.environ.get("WANDB_RUN_GROUP"))
+        return WandbLogger(project=project, name=run, group=group, resume=resume, **kwargs)
     if logger_name == "mlflow":
         return MLFlowLogger(experiment_name=name, **kwargs)
     raise ValueError(f"`--logger_name={logger_name}` is not a valid option. Choose from 'csv', 'tensorboard', 'wandb'.")