fix dump sl

Author: YassineYousfi

run_train.sh

@@ -38,7 +38,6 @@ CONFIG=${CONFIG:-"llama3_debugmodel"}
 COMM_MODE=${COMM_MODE:-""}
 NNODES=${NNODES:-${SLURM_JOB_NUM_NODES:-1}}
 NODE_RANK=${NODE_RANK:-${SLURM_NODEID:-0}}
-TORCHTITAN_DUMP_FOLDER=${TORCHTITAN_DUMP_FOLDER:-"/tmp/torchtitan_train/"}
 
 # xx injects these
 TORCHTITAN_ARGS=()
@@ -88,7 +87,7 @@ export REPORTERV2_TRAINING_ID
 if [ -n "$COMM_MODE" ]; then
     # Communication mode specified: validate configuration or run in debug mode
     echo "Running with comm_mode=${COMM_MODE}"
-    NGPU="${NGPU}" LOCAL_RANK=0 python3 -m torchtitan.train --module ${MODULE} --config ${CONFIG} --dump-folder "${TORCHTITAN_DUMP_FOLDER}" "$@" --comm.mode=${COMM_MODE} --training.steps 1
+    NGPU="${NGPU}" LOCAL_RANK=0 python3 -m torchtitan.train --module ${MODULE} --config ${CONFIG} "$@" --comm.mode=${COMM_MODE} --training.steps 1
 else
     if [[ -n "${MASTER_ADDR:-}" ]]; then
         RDZV_ENDPOINT="${MASTER_ADDR}:${MASTER_PORT}"
@@ -103,5 +102,5 @@ else
     --rdzv_id=${RDZV_ID:-${SLURM_JOB_ID:-$(generate_uuid)}} --rdzv_backend c10d \
     --rdzv_endpoint="${RDZV_ENDPOINT}" \
     --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
-    -m torchtitan.train --module ${MODULE} --config ${CONFIG} --dump-folder "${TORCHTITAN_DUMP_FOLDER}" "$@"
+    -m torchtitan.train --module ${MODULE} --config ${CONFIG} "$@"
 fi

tests/unit_tests/test_checkpoint.py

@@ -386,10 +386,11 @@ def test_dcp_save_load_memory_fs_checkpoint(self):
             states=self.states,
             config=cfg,
             sd_adapter=None,
-            base_folder="",
+            base_folder="./outputs",
         )
 
         try:
+            self.assertEqual(manager.folder, root)
             checkpoint_id = manager._create_checkpoint_id(1)
             expected = torch.tensor([3.0])
             manager.dcp_save(

torchtitan/components/checkpoint.py

@@ -16,6 +16,7 @@
 from typing import Any, cast, Literal
 from urllib.parse import urlparse
 
+from fsspec.core import split_protocol
 import torch
 import torch.distributed as dist
 import torch.distributed.checkpoint as dcp
@@ -458,7 +459,11 @@ def __init__(
         if not self.enable:
             return
 
-        self.folder = fs.join_path(base_folder, config.folder)
+        self.folder = (
+            config.folder
+            if split_protocol(config.folder)[0] is not None
+            else fs.join_path(base_folder, config.folder)
+        )
         self.checkpoint_id_format = config.checkpoint_id_format
         self.interval = config.interval
 

torchtitan/experiments/worldmodel/config_registry.py

@@ -63,11 +63,10 @@ def worldmodel() -> WorldModelTrainer.Config:
     optimizer = default_adamw(lr=2e-4, weight_decay=1e-2)
     optimizer.implementation = "fused_opt_states_bf16"
     local_world_size, world_size, num_nodes = _world_sizes()
-    checkpoint_base_folder = _reporterv2_checkpoint_base_folder()
+    checkpoint_folder = _reporterv2_checkpoint_folder()
 
     return WorldModelTrainer.Config(
         hf_assets_path=".",
-        dump_folder=checkpoint_base_folder or "./outputs",
         loss=WorldModelLoss.Config(plan_loss_weight=0.1),
         tokenizer=WorldModelTokenizer.Config(
             compressor_model=COMPRESSOR_MODEL,
@@ -111,7 +110,7 @@ def worldmodel() -> WorldModelTrainer.Config:
         ),
         checkpoint=WorldModelTorchPackageCheckpointManager.Config(
             enable=True,
-            folder=os.getenv("REPORTERV2_TRAINING_ID") or "checkpoint",
+            folder=checkpoint_folder,
             interval=validation_freq * 5,
             async_mode="async",
             keep_latest_k=0,
@@ -155,9 +154,10 @@ def _world_sizes() -> tuple[int, int, int]:
     return local_world_size, world_size, num_nodes
 
 
-def _reporterv2_checkpoint_base_folder() -> str:
+def _reporterv2_checkpoint_folder() -> str:
     host = os.getenv("REPORTERV2_HOST")
-    return f"{host.rstrip('/')}/checkpoint" if host else ""
+    folder = os.getenv("REPORTERV2_TRAINING_ID") or "checkpoint"
+    return f"{host.rstrip('/')}/checkpoint/{folder}" if host else folder
 
 
 def main() -> None:

torchtitan/experiments/worldmodel/torchpackage_checkpoint.py

@@ -36,7 +36,8 @@
 PACKAGE_NAME = "model.torchpackage"
 MODEL_CONFIG_FILE = "_torchpackage_model_config.pt"
 STRUCTURED_LOG_DIR = os.getenv(
-    "TORCHTITAN_STRUCTURED_LOG_DIR", "./outputs/worldmodel_torchpackage_checkpoint"
+    "TORCHTITAN_STRUCTURED_LOG_DIR",
+    "/tmp/torchtitan_train/worldmodel_torchpackage_checkpoint",
 )
 WORLD_MODEL_TORCH_PACKAGE_RECIPE = (
     "torchtitan.experiments.worldmodel.torchpackage_checkpoint:"

torchtitan/trainer.py

@@ -73,7 +73,7 @@ class Config(Configurable.Config):
         (fqn to file mapping), the config.json file, generation_config.json, and tokenizer files.
         """
 
-        dump_folder: str = "./outputs"
+        dump_folder: str = "/tmp/torchtitan_train"
         """Folder to dump job outputs"""
 
         codedir: str = ""