diff --git a/apps/grpo/llama3_8b.yaml b/apps/grpo/llama3_8b.yaml index 0597bb238..b989ed04d 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: false # Global configuration group_size: 4 diff --git a/apps/grpo/qwen3_1_7b.yaml b/apps/grpo/qwen3_1_7b.yaml index ebdd27787..5410fa707 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: false + # 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..f49466b7f 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: 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 6736b0b01..235d889bb 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: 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 d9a5a9783..d44a9e07b 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: false # 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..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 @@ -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) diff --git a/apps/sft/qwen3_8b.yaml b/apps/sft/qwen3_8b.yaml index a8d2244a9..27f2f43b5 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: false # 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, )