From 7e8c8ecb11dd70e9c89efc5e47d986b26c21d5fe Mon Sep 17 00:00:00 2001 From: rakkit <26144573+rakkit@users.noreply.github.com> Date: Tue, 9 Dec 2025 16:54:27 +0100 Subject: [PATCH 1/3] allow checkpoint and logger to set base dir --- apps/grpo/llama3_8b.yaml | 3 +++ apps/grpo/qwen3_1_7b.yaml | 4 ++++ apps/grpo/qwen3_32b.yaml | 4 +++- apps/grpo/qwen3_8b.yaml | 4 +++- apps/sft/llama3_8b.yaml | 3 +++ apps/sft/main.py | 7 +++++-- apps/sft/qwen3_8b.yaml | 4 +++- src/forge/observability/metric_actors.py | 5 ++++- src/forge/observability/metrics.py | 21 ++++++++++++++++++--- 9 files changed, 46 insertions(+), 9 deletions(-) diff --git a/apps/grpo/llama3_8b.yaml b/apps/grpo/llama3_8b.yaml index 0597bb238..29974c898 100644 --- a/apps/grpo/llama3_8b.yaml +++ b/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: true # Global configuration group_size: 4 diff --git a/apps/grpo/qwen3_1_7b.yaml b/apps/grpo/qwen3_1_7b.yaml index ebdd27787..17a93ea05 100644 --- a/apps/grpo/qwen3_1_7b.yaml +++ b/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: true + # Global configuration group_size: 8 local_batch_size: 16 # per-device batch size diff --git a/apps/grpo/qwen3_32b.yaml b/apps/grpo/qwen3_32b.yaml index 0366b58a3..4441039d7 100644 --- a/apps/grpo/qwen3_32b.yaml +++ b/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: true # Global configuration group_size: 16 local_batch_size: 32 # per-device batch size diff --git a/apps/grpo/qwen3_8b.yaml b/apps/grpo/qwen3_8b.yaml index 6736b0b01..f710665c5 100644 --- a/apps/grpo/qwen3_8b.yaml +++ b/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: true # Global configuration group_size: 16 local_batch_size: 4 # per-device batch size diff --git a/apps/sft/llama3_8b.yaml b/apps/sft/llama3_8b.yaml index d9a5a9783..4635f2a77 100644 --- a/apps/sft/llama3_8b.yaml +++ b/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: true # TODO: required by torchtitan # https://github.com/pytorch/torchtitan/blob/2f1c814da071cc8ad165d00be6f9c1a66f8e1cce/torchtitan/distributed/utils.py#L265 diff --git a/apps/sft/main.py b/apps/sft/main.py index 4f2a7be74..1883bed04 100644 --- a/apps/sft/main.py +++ b/apps/sft/main.py @@ -65,7 +65,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 @@ -140,7 +140,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 +480,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) diff --git a/apps/sft/qwen3_8b.yaml b/apps/sft/qwen3_8b.yaml index a8d2244a9..015331044 100644 --- a/apps/sft/qwen3_8b.yaml +++ b/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: true # TODO: required by torchtitan # https://github.com/pytorch/torchtitan/blob/2f1c814da071cc8ad165d00be6f9c1a66f8e1cce/torchtitan/distributed/utils.py#L265 comm: diff --git a/src/forge/observability/metric_actors.py b/src/forge/observability/metric_actors.py index f7b1ccedf..8255f7917 100644 --- a/src/forge/observability/metric_actors.py +++ b/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, diff --git a/src/forge/observability/metrics.py b/src/forge/observability/metrics.py index 2c5a4663f..531b06230 100644 --- a/src/forge/observability/metrics.py +++ b/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, ) From 55ad01f874b8021dcd9d33aefab9e68b22b6fb41 Mon Sep 17 00:00:00 2001 From: rakkit <26144573+rakkit@users.noreply.github.com> Date: Tue, 9 Dec 2025 17:07:45 +0100 Subject: [PATCH 2/3] by default dont pring config --- apps/grpo/llama3_8b.yaml | 2 +- apps/grpo/qwen3_1_7b.yaml | 2 +- apps/grpo/qwen3_32b.yaml | 2 +- apps/grpo/qwen3_8b.yaml | 2 +- apps/sft/llama3_8b.yaml | 2 +- apps/sft/qwen3_8b.yaml | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/apps/grpo/llama3_8b.yaml b/apps/grpo/llama3_8b.yaml index 29974c898..b989ed04d 100644 --- a/apps/grpo/llama3_8b.yaml +++ b/apps/grpo/llama3_8b.yaml @@ -2,7 +2,7 @@ # >>> python -m apps.grpo.main --config apps/grpo/llama3_8b.yaml job: dump_folder: ./ - print_config: true + print_config: false # Global configuration group_size: 4 diff --git a/apps/grpo/qwen3_1_7b.yaml b/apps/grpo/qwen3_1_7b.yaml index 17a93ea05..5410fa707 100644 --- a/apps/grpo/qwen3_1_7b.yaml +++ b/apps/grpo/qwen3_1_7b.yaml @@ -3,7 +3,7 @@ job: dump_folder: ./ - print_config: true + print_config: false # Global configuration group_size: 8 diff --git a/apps/grpo/qwen3_32b.yaml b/apps/grpo/qwen3_32b.yaml index 4441039d7..f49466b7f 100644 --- a/apps/grpo/qwen3_32b.yaml +++ b/apps/grpo/qwen3_32b.yaml @@ -3,7 +3,7 @@ # NOTE - This has not been tested for correctness yet! All testing so far has been only for infrastructure stability job: dump_folder: ./ - print_config: true + print_config: false # Global configuration group_size: 16 local_batch_size: 32 # per-device batch size diff --git a/apps/grpo/qwen3_8b.yaml b/apps/grpo/qwen3_8b.yaml index f710665c5..235d889bb 100644 --- a/apps/grpo/qwen3_8b.yaml +++ b/apps/grpo/qwen3_8b.yaml @@ -2,7 +2,7 @@ # >>> python -m apps.grpo.main --config apps/grpo/qwen3_8b.yaml job: dump_folder: ./ - print_config: true + print_config: false # Global configuration group_size: 16 local_batch_size: 4 # per-device batch size diff --git a/apps/sft/llama3_8b.yaml b/apps/sft/llama3_8b.yaml index 4635f2a77..d44a9e07b 100644 --- a/apps/sft/llama3_8b.yaml +++ b/apps/sft/llama3_8b.yaml @@ -3,7 +3,7 @@ # Config for supervised full finetuning using a Llama3.1 8B Instruct model job: dump_folder: ./ - print_config: true + print_config: false # TODO: required by torchtitan # https://github.com/pytorch/torchtitan/blob/2f1c814da071cc8ad165d00be6f9c1a66f8e1cce/torchtitan/distributed/utils.py#L265 diff --git a/apps/sft/qwen3_8b.yaml b/apps/sft/qwen3_8b.yaml index 015331044..27f2f43b5 100644 --- a/apps/sft/qwen3_8b.yaml +++ b/apps/sft/qwen3_8b.yaml @@ -2,7 +2,7 @@ job: dump_folder: ./ - print_config: true + print_config: false # TODO: required by torchtitan # https://github.com/pytorch/torchtitan/blob/2f1c814da071cc8ad165d00be6f9c1a66f8e1cce/torchtitan/distributed/utils.py#L265 comm: From bcd48750c4e4ddced46951a7b5a7d572fd5f7bf0 Mon Sep 17 00:00:00 2001 From: rakkit <26144573+rakkit@users.noreply.github.com> Date: Tue, 9 Dec 2025 20:21:40 +0100 Subject: [PATCH 3/3] print config when its possible for sft --- apps/sft/main.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/apps/sft/main.py b/apps/sft/main.py index 1883bed04..948d26cdb 100644 --- a/apps/sft/main.py +++ b/apps/sft/main.py @@ -12,6 +12,7 @@ import asyncio import contextlib +import json import logging import math import os @@ -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"]