Move config into utils (#453)

Author: Jack-Khuu

.meta/mast/main.py

@@ -9,7 +9,6 @@
 import sys
 
 from apps.grpo.main import main as grpo_main
-from forge.cli.config import parse
 from forge.controller.launcher import (
     JOB_NAME_KEY,
     LAUNCHER_KEY,
@@ -25,6 +24,7 @@
     ProvisionerConfig,
     ServiceConfig,
 )
+from forge.util.config import parse
 from omegaconf import DictConfig
 
 DEFAULT_CHECKPOINT_FOLDER_KEY = "checkpoint_folder"

apps/grpo/main.py

@@ -24,7 +24,6 @@
 from forge.actors.reference_model import ReferenceModel
 from forge.actors.replay_buffer import ReplayBuffer
 from forge.actors.trainer import RLTrainer
-from forge.cli.config import parse
 from forge.controller.actor import ForgeActor
 from forge.controller.provisioner import init_provisioner, shutdown
 from forge.data.rewards import MathReward, ThinkingReward
@@ -34,6 +33,7 @@
 from forge.observability.perf_tracker import Tracer
 
 from forge.types import LauncherConfig, ProvisionerConfig
+from forge.util.config import parse
 from forge.util.ops import compute_logprobs
 from monarch.actor import endpoint
 from omegaconf import DictConfig

apps/sft/main.py

@@ -22,12 +22,12 @@
 import torch
 
 import torchtitan.experiments.forge.train_spec as forge_train_spec
-from forge.cli.config import parse
 from forge.controller import ForgeActor
 from forge.data.collate import collate_packed
 from forge.data.datasets.packed import PackedDataset, TextPacker
 from forge.data.datasets.sft_dataset import AlpacaToMessages, sft_iterable_dataset
 from forge.data.tokenizer import HuggingFaceModelTokenizer
+from forge.util.config import parse
 
 from monarch.actor import current_rank, current_size, endpoint
 from omegaconf import DictConfig, OmegaConf

src/forge/cli/__init__.py

@@ -17,11 +17,11 @@
 from forge.actors.generator import Generator
 
 from forge.actors.trainer import RLTrainer
-from forge.cli.config import resolve_hf_hub_paths
 from forge.controller.provisioner import init_provisioner
 
 from forge.controller.service.service import uuid
 from forge.types import LauncherConfig, ProvisionerConfig
+from forge.util.config import resolve_hf_hub_paths
 from monarch.actor import endpoint
 
 from omegaconf import DictConfig, OmegaConf

src/forge/cli/config.py renamed to src/forge/util/config.py

@@ -25,9 +25,9 @@
 import torch
 
 from forge.actors.reference_model import ReferenceModel
-from forge.cli.config import _resolve_hf_model_path
 from forge.controller import ForgeActor
 from forge.controller.provisioner import shutdown
+from forge.util.config import _resolve_hf_model_path
 from monarch.actor import endpoint
 from torchtitan.config.job_config import Checkpoint, Compile, Model, Parallelism
 from transformers import AutoModelForCausalLM, AutoTokenizer

tests/integration_tests/test_policy_update.py

@@ -11,7 +11,6 @@
 import torch
 import torchstore as ts
 from forge.actors.trainer import RLTrainer
-from forge.cli.config import parse
 from forge.controller.launcher import JOB_NAME_KEY, LAUNCHER_KEY
 from forge.controller.provisioner import init_provisioner, shutdown
 from forge.observability.metric_actors import get_or_create_metric_logger
@@ -23,6 +22,7 @@
     ProvisionerConfig,
     ServiceConfig,
 )
+from forge.util.config import parse
 from omegaconf import DictConfig
 from vllm.transformers_utils.tokenizer import get_tokenizer
 

tests/integration_tests/test_titan_fwd_vs_hf_fwd.py

@@ -18,13 +18,13 @@
 from forge.actors._torchstore_utils import get_param_key
 from forge.actors.generator import Generator
 from forge.actors.replay_buffer import ReplayBuffer
-from forge.cli.config import parse
 from forge.controller.actor import ForgeActor
 from forge.controller.provisioner import shutdown
 from forge.losses.grpo_loss import SimpleGRPOLoss
 from forge.observability.metric_actors import get_or_create_metric_logger
 
 from forge.observability.metrics import record_metric, Reduce
+from forge.util.config import parse
 from forge.util.ops import selective_log_softmax
 from monarch.actor import endpoint
 from omegaconf import DictConfig

tests/sandbox/rl_trainer/main.py

@@ -14,13 +14,13 @@
 import os
 
 from forge.actors.generator import Generator
-from forge.cli.config import parse
 
 from forge.controller.provisioner import init_provisioner, shutdown
 
 from forge.data_models.completion import Completion
 from forge.observability.metric_actors import get_or_create_metric_logger
 from forge.types import LauncherConfig, ProvisionerConfig
+from forge.util.config import parse
 from omegaconf import DictConfig
 
 os.environ["HYPERACTOR_MESSAGE_DELIVERY_TIMEOUT_SECS"] = "600"