Standardize out_dir behavior (#1338)

Author: rasbt

litgpt/finetune/adapter.py

@@ -29,6 +29,7 @@
     chunked_cross_entropy,
     copy_config_files,
     get_default_supported_precision,
+    init_out_dir,
     load_checkpoint,
     num_parameters,
     parse_devices,
@@ -61,7 +62,8 @@ def setup(
 
     Arguments:
         checkpoint_dir: The path to the base model's checkpoint directory to load for finetuning.
-        out_dir: Directory in which to save checkpoints and logs.
+        out_dir: Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
+            /teamspace/jobs/<job-name>/share.
         precision: The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true".
         quantize: If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information.
         devices: How many devices/GPUs to use.
@@ -75,6 +77,7 @@ def setup(
     pprint(locals())
     data = Alpaca() if data is None else data
     devices = parse_devices(devices)
+    out_dir = init_out_dir(out_dir)
 
     check_valid_checkpoint_dir(checkpoint_dir)
     config = Config.from_file(checkpoint_dir / "model_config.yaml")

litgpt/finetune/adapter_v2.py

@@ -29,6 +29,7 @@
     chunked_cross_entropy,
     copy_config_files,
     get_default_supported_precision,
+    init_out_dir,
     load_checkpoint,
     num_parameters,
     parse_devices,
@@ -61,7 +62,8 @@ def setup(
 
     Arguments:
         checkpoint_dir: The path to the base model's checkpoint directory to load for finetuning.
-        out_dir: Directory in which to save checkpoints and logs.
+        out_dir: Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
+            /teamspace/jobs/<job-name>/share.
         precision: The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true".
         quantize: If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information.
         devices: How many devices/GPUs to use.
@@ -75,6 +77,7 @@ def setup(
     pprint(locals())
     data = Alpaca() if data is None else data
     devices = parse_devices(devices)
+    out_dir = init_out_dir(out_dir)
 
     check_valid_checkpoint_dir(checkpoint_dir)
     config = Config.from_file(checkpoint_dir / "model_config.yaml")

litgpt/finetune/full.py

@@ -28,6 +28,7 @@
     copy_config_files,
     get_default_supported_precision,
     load_checkpoint,
+    init_out_dir,
     num_parameters,
     parse_devices,
     save_hyperparameters,
@@ -59,7 +60,8 @@ def setup(
 
     Arguments:
         checkpoint_dir: The path to the base model's checkpoint directory to load for finetuning.
-        out_dir: Directory in which to save checkpoints and logs.
+        out_dir: Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
+            /teamspace/jobs/<job-name>/share.
         precision: The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true".
         devices: How many devices/GPUs to use
         resume: Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
@@ -74,6 +76,7 @@ def setup(
     pprint(locals())
     data = Alpaca() if data is None else data
     devices = parse_devices(devices)
+    out_dir = init_out_dir(out_dir)
 
     check_valid_checkpoint_dir(checkpoint_dir)
     config = Config.from_file(checkpoint_dir / "model_config.yaml")

litgpt/finetune/lora.py

@@ -31,6 +31,7 @@
     copy_config_files,
     get_default_supported_precision,
     load_checkpoint,
+    init_out_dir,
     num_parameters,
     parse_devices,
     save_hyperparameters,
@@ -71,7 +72,8 @@ def setup(
 
     Arguments:
         checkpoint_dir: The path to the base model's checkpoint directory to load for finetuning.
-        out_dir: Directory in which to save checkpoints and logs.
+        out_dir: Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
+            /teamspace/jobs/<job-name>/share.
         precision: The precision to use for finetuning. Possible choices: "bf16-true", "bf16-mixed", "32-true".
         quantize: If set, quantize the model with this algorithm. See ``tutorials/quantize.md`` for more information.
         devices: How many devices/GPUs to use.
@@ -94,6 +96,7 @@ def setup(
     pprint(locals())
     data = Alpaca() if data is None else data
     devices = parse_devices(devices)
+    out_dir = init_out_dir(out_dir)
 
     check_valid_checkpoint_dir(checkpoint_dir)
     config = Config.from_file(

litgpt/pretrain.py

@@ -1,7 +1,6 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
 import math
-import os
 import pprint
 import time
 from datetime import timedelta
@@ -30,6 +29,7 @@
     choose_logger,
     chunked_cross_entropy,
     copy_config_files,
+    init_out_dir,
     num_parameters,
     parse_devices,
     reset_parameters,
@@ -404,12 +404,6 @@ def init_weights(module, std):
         reset_parameters(model)
 
 
-def init_out_dir(out_dir: Path) -> Path:
-    if not out_dir.is_absolute() and "LIGHTNING_ARTIFACTS_DIR" in os.environ:
-        return Path(os.getenv("LIGHTNING_ARTIFACTS_DIR")) / out_dir
-    return out_dir
-
-
 def save_checkpoint(fabric, state, tokenizer_dir, checkpoint_file):
     model = state["model"]
     checkpoint_file.parent.mkdir(parents=True, exist_ok=True)

litgpt/utils.py

@@ -3,6 +3,7 @@
 """Utility functions for training and inference."""
 import inspect
 import math
+import os
 import pickle
 import shutil
 import sys
@@ -27,6 +28,12 @@
     from litgpt import GPT, Config
 
 
+def init_out_dir(out_dir: Path) -> Path:
+    if not out_dir.is_absolute() and "LIGHTNING_ARTIFACTS_DIR" in os.environ:
+        return Path(os.getenv("LIGHTNING_ARTIFACTS_DIR")) / out_dir
+    return out_dir
+
+
 def find_multiple(n: int, k: int) -> int:
     assert k > 0
     if n % k == 0:

tests/test_pretrain.py

@@ -13,10 +13,11 @@
 from lightning.fabric.strategies import FSDPStrategy, SingleDeviceStrategy
 from torch.utils.data import DataLoader
 
+from test_utils import test_init_out_dir
 from litgpt import pretrain
 from litgpt.args import EvalArgs, TrainArgs
 from litgpt.config import Config
-from litgpt.pretrain import init_out_dir, initialize_weights
+from litgpt.pretrain import initialize_weights
 
 
 @RunIf(min_cuda_gpus=2, standalone=True)
@@ -89,17 +90,6 @@ def test_pretrain_model_name_and_config():
         pretrain.setup(model_name="tiny-llama-1.1b", model_config=Config(name="tiny-llama-1.1b"))
 
 
-def test_init_out_dir(tmp_path):
-    relative_path = Path("./out")
-    absolute_path = tmp_path / "out"
-    assert init_out_dir(relative_path) == relative_path
-    assert init_out_dir(absolute_path) == absolute_path
-
-    with mock.patch.dict(os.environ, {"LIGHTNING_ARTIFACTS_DIR": "prefix"}):
-        assert init_out_dir(relative_path) == Path("prefix") / relative_path
-        assert init_out_dir(absolute_path) == absolute_path
-
-
 @pytest.mark.parametrize(("strategy", "expected"), [(SingleDeviceStrategy, True), (FSDPStrategy, False)])
 def test_initialize_weights(strategy, expected):
     fabric_mock = Mock()

tests/test_utils.py

@@ -30,6 +30,7 @@
     copy_config_files,
     find_multiple,
     incremental_save,
+    init_out_dir,
     num_parameters,
     parse_devices,
     save_hyperparameters,
@@ -294,3 +295,14 @@ def test_choose_logger(tmp_path):
 
     with pytest.raises(ValueError, match="`--logger_name=foo` is not a valid option."):
         choose_logger("foo", out_dir=tmp_path, name="foo")
+
+
+def test_init_out_dir(tmp_path):
+    relative_path = Path("./out")
+    absolute_path = tmp_path / "out"
+    assert init_out_dir(relative_path) == relative_path
+    assert init_out_dir(absolute_path) == absolute_path
+
+    with mock.patch.dict(os.environ, {"LIGHTNING_ARTIFACTS_DIR": "prefix"}):
+        assert init_out_dir(relative_path) == Path("prefix") / relative_path
+        assert init_out_dir(absolute_path) == absolute_path