allow set base_dir for checkpoint and logger

apps/grpo/llama3_8b.yaml

@@ -1,5 +1,8 @@
 # Grouped Relative Policy Optimization (GRPO)
 # >>> python -m apps.grpo.main --config apps/grpo/llama3_8b.yaml
+job:
+  dump_folder: ./
+  print_config: false
 
 # Global configuration
 group_size: 4

apps/grpo/qwen3_1_7b.yaml

@@ -1,6 +1,10 @@
 # Grouped Relative Policy Optimization (GRPO)
 # >>> python -m apps.grpo.main --config apps/grpo/qwen3_1_7b.yaml
 
+job:
+  dump_folder: ./
+  print_config: false
+
 # Global configuration
 group_size: 8
 local_batch_size: 16 # per-device batch size

apps/grpo/qwen3_32b.yaml

@@ -1,7 +1,9 @@
 # Grouped Relative Policy Optimization (GRPO)
 # >>> python -m apps.grpo.main --config apps/grpo/qwen3_32b.yaml
 # NOTE - This has not been tested for correctness yet! All testing so far has been only for infrastructure stability
-
+job:
+  dump_folder: ./
+  print_config: false
 # Global configuration
 group_size: 16
 local_batch_size: 32 # per-device batch size

apps/grpo/qwen3_8b.yaml

@@ -1,6 +1,8 @@
 # Grouped Relative Policy Optimization (GRPO)
 # >>> python -m apps.grpo.main --config apps/grpo/qwen3_8b.yaml
-
+job:
+  dump_folder: ./
+  print_config: false
 # Global configuration
 group_size: 16
 local_batch_size: 4 # per-device batch size

apps/sft/llama3_8b.yaml

@@ -1,6 +1,9 @@
 # >>> python -m apps.sft.main --config apps/sft/llama3_8b.yaml
 
 # Config for supervised full finetuning using a Llama3.1 8B Instruct model
+job:
+  dump_folder: ./
+  print_config: false
 
 # TODO: required by torchtitan
 # https://github.com/pytorch/torchtitan/blob/2f1c814da071cc8ad165d00be6f9c1a66f8e1cce/torchtitan/distributed/utils.py#L265

apps/sft/main.py

@@ -12,6 +12,7 @@
 
 import asyncio
 import contextlib
+import json
 import logging
 import math
 import os
@@ -65,7 +66,7 @@ class ForgeSFTRecipe(ForgeActor, ForgeEngine):
     tokenizer: Tokenizer
     train_dataloader: Dataloader
     # val_dataloader: Dataloader
-    metric_logger: MetricLogger
+    mlogger: MetricLogger
     profiler: Profiler
     device: torch.device
     step: int
@@ -95,6 +96,12 @@ def record_batch_metrics(self, data_metrics: list):
 
     @endpoint
     async def setup(self):
+        # print config
+        if self.job_config.get("job", {}).get("print_config", False):
+            dict_config = OmegaConf.to_container(self.job_config, resolve=True)
+            formatted_config = json.dumps(dict_config, indent=2, ensure_ascii=False)
+            logger.info(f"Running with configs: {formatted_config}")
+
         # Validate that compile is only used with flex attention
         if self.job_config.training.compile:
             raise ValueError(
@@ -116,7 +123,7 @@ async def setup(self):
         train_datasets_config = self.job_config.training.datasets
         self.train_dataloader = self.setup_data(train_datasets_config)
 
-        # Load eval datasets
+        # Load eval datasetsf
         eval_config = self.job_config["eval"]
         self.val_dataloaders = {}
         self.eval_every_n_steps = eval_config["eval_every_n_steps"]
@@ -140,7 +147,9 @@ async def setup(self):
 
         # TODO: confirm that this is working properly
         # Should also use load, not dcp_load
-        self.checkpointer.load(step=self.current_step)
+        if self.current_step != 0:
+            # should skip load if current_step is 0
+            self.checkpointer.load(step=self.current_step)
 
         # self.profiler = self.setup_profiler(self.train_config.profiler_config)
         # self.logger = self.setup_logger(self.train_config.logger_config)
@@ -478,6 +487,7 @@ async def run(cfg: DictConfig) -> None:
 
     # Initialize metric logger in main process
     metric_logging_cfg = cfg.get("metric_logging", {})
+    metric_logging_cfg["dump_folder"] = cfg.get("job").get("dump_folder", "./")
     mlogger = await get_or_create_metric_logger(process_name="Controller")
     await mlogger.init_backends.call_one(metric_logging_cfg)
 

apps/sft/qwen3_8b.yaml

@@ -1,6 +1,8 @@
 # >>> python -m apps.sft.main --config apps/sft/qwen3_8b.yaml
 
-
+job:
+  dump_folder: ./
+  print_config: false
 # TODO: required by torchtitan
 # https://github.com/pytorch/torchtitan/blob/2f1c814da071cc8ad165d00be6f9c1a66f8e1cce/torchtitan/distributed/utils.py#L265
 comm:

src/forge/observability/metric_actors.py

@@ -309,6 +309,8 @@ async def init_backends(
         self.config = {}
         self.run_config = run_config
 
+        base_folder = backend_config.pop("dump_folder", "./")
+
         # Skip initialization if disabled by environment flag
         if FORGE_DISABLE_METRICS.get_value():
             return
@@ -322,7 +324,8 @@ async def init_backends(
             mode = backend_config["logging_mode"]
 
             backend: LoggerBackend = get_logger_backend_class(backend_name)(
-                **backend_config
+                **backend_config,
+                base_folder=base_folder,
             )
             await backend.init(
                 role=BackendRole.GLOBAL,

src/forge/observability/metrics.py

@@ -769,10 +769,16 @@ class LoggerBackend(ABC):
     """
 
     def __init__(
-        self, *, logging_mode: LoggingMode, per_rank_share_run: bool = False, **kwargs
+        self,
+        *,
+        logging_mode: LoggingMode,
+        per_rank_share_run: bool = False,
+        base_folder: str = "./",
+        **kwargs,
     ) -> None:
         self.logging_mode = logging_mode
         self.per_rank_share_run = per_rank_share_run
+        self.base_folder = base_folder
         self.backend_kwargs = kwargs
 
     @abstractmethod
@@ -933,6 +939,7 @@ async def init(
         controller_logger_metadata: dict[str, Any] | None = None,
         process_name: str | None = None,
         run_config: dict[str, Any] | None = None,
+        base_folder: str | None = None,
     ) -> None:
         if controller_logger_metadata is None:
             controller_logger_metadata = {}
@@ -972,14 +979,20 @@ async def _init_global(self, run_name: str | None):
         import wandb
 
         self.run = wandb.init(
-            name=run_name, config=self.run_config, **self.backend_kwargs
+            name=run_name,
+            config=self.run_config,
+            dir=self.base_folder,
+            **self.backend_kwargs,
         )
 
     async def _init_per_rank(self, run_name: str):
         import wandb
 
         self.run = wandb.init(
-            name=run_name, config=self.run_config, **self.backend_kwargs
+            name=run_name,
+            config=self.run_config,
+            dir=self.base_folder,
+            **self.backend_kwargs,
         )
 
     async def _init_shared_global(self, run_name: str | None):
@@ -991,6 +1004,7 @@ async def _init_shared_global(self, run_name: str | None):
         self.run = wandb.init(
             name=run_name,
             config=self.run_config,
+            dir=self.base_folder,
             settings=settings,
             **self.backend_kwargs,
         )
@@ -1014,6 +1028,7 @@ async def _init_shared_local(
             id=shared_id,
             config=self.run_config,
             settings=settings,
+            dir=self.base_folder,
             **self.backend_kwargs,
         )