move the setting of PYTORCH_CUDA_ALLOC_CONF to the cli rather than tr…

…ain module

src/axolotl/cli/main.py

@@ -1,4 +1,6 @@
 """CLI definition for various axolotl commands."""
+import os
+
 # pylint: disable=redefined-outer-name
 import subprocess  # nosec B404
 from typing import Optional
@@ -14,6 +16,7 @@
 )
 from axolotl.common.cli import EvaluateCliArgs, PreprocessCliArgs, TrainerCliArgs
 from axolotl.utils.config.models.input.v0_4_1 import AxolotlInputConfig
+from axolotl.utils.trainer import set_pytorch_cuda_alloc_conf
 
 
 @click.group()
@@ -48,6 +51,9 @@ def train(config: str, accelerate: bool, **kwargs):
     """Train or fine-tune a model."""
     kwargs = {k: v for k, v in kwargs.items() if v is not None}
 
+    # Enable expandable segments for cuda allocation to improve VRAM usage
+    set_pytorch_cuda_alloc_conf()
+
     if accelerate:
         base_cmd = ["accelerate", "launch", "-m", "axolotl.cli.train"]
         if config:

src/axolotl/train.py

@@ -24,7 +24,7 @@
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.freeze import freeze_layers_except
 from axolotl.utils.models import load_model, load_processor, load_tokenizer
-from axolotl.utils.trainer import set_pytorch_cuda_alloc_conf, setup_trainer
+from axolotl.utils.trainer import setup_trainer
 
 try:
     from optimum.bettertransformer import BetterTransformer
@@ -53,9 +53,6 @@ class TrainDatasetMeta:
 def train(
     *, cfg: DictDefault, cli_args: TrainerCliArgs, dataset_meta: TrainDatasetMeta
 ) -> Tuple[Union[PeftModel, PreTrainedModel], PreTrainedTokenizer]:
-    # Enable expandable segments for cuda allocation to improve VRAM usage
-    set_pytorch_cuda_alloc_conf()
-
     # Load tokenizer
     LOG.debug(
         f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}",