chatlearn/utils/arguments.py

# Copyright 2024 Alibaba Group Holding Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""arguments from command or yaml."""

import argparse
import ast
import os
from typing import List, Optional, Union
import re

import yaml

from chatlearn.utils.constant import (
    DYNAMIC_BATCH_SIZE, LORA_LAYER, RAY_PG_STRATEGY,
    PARAM_SYNC_COMM_TYPE, ROUTED_EXPERT_REGROUPING_COMM_TYPE,
    TrainingShffuleMode)
from chatlearn.utils.logger import logger
from chatlearn.utils.utils import get_attributes


def get_path(fn, folder):
    if not fn.startswith("/") and not fn.startswith(folder):
        fn = os.path.join(folder, fn)
    assert os.path.exists(fn), f'{fn} not exists'
    return fn


def convert_type(data):
    try:
        return ast.literal_eval(data)
    except Exception:
        return data


def parse_value(value):
    if isinstance(value, dict):
        return {k: parse_value(v) for k, v in value.items()}

    if isinstance(value, str):
        if value.strip().startswith("${"):
            # ${env_name:default_value}
            placeholder = value.replace("${", "")[:-1]
            placeholder = placeholder.split(":")
            env_name = placeholder[0]
            if env_name in os.environ:
                value = convert_type(os.environ[env_name])
            else:
                if len(placeholder) > 1:
                    value = convert_type(placeholder[1])
                else:
                    logger.warning(f"cannot find value for {env_name}, set to None")
                    value = None
        # handling scientific notation(e.g., "5e-6", "5E+10")
        elif re.match(r"^[+-]?(\d+\.?\d*|\.\d+)([eE][+-]?\d+)?$", value):
            try:
                value = float(value)
            except Exception:
                pass
    return value


def update_dict(src, dst):
    # do not overwrite
    for k, v in src.items():
        if k not in dst:
            dst[k] = v
        else:
            if isinstance(v, dict) and isinstance(dst[k], dict):
                update_dict(v, dst[k])


def parse_args_from_yaml(config_file, config_dir):
    with open(config_file, 'r', encoding='utf-8') as stream:
        config_vars = yaml.load(stream, Loader=yaml.SafeLoader)
        # empty yaml file
        if config_vars is None:
            return {}
        config_vars = {key: parse_value(value) for key, value in config_vars.items()}
        if 'includes' in config_vars:
            includes_vars = {}
            # iterate in reverse order, so the next include overwrite the prev
            for base in reversed(config_vars["includes"]):
                base_path = get_path(base, config_dir)
                base_config = parse_args_from_yaml(base_path, config_dir)
                update_dict(base_config, includes_vars)
            update_dict(includes_vars, config_vars)
        return config_vars


def parse_args():
    """Parse all arguments."""
    parser = argparse.ArgumentParser(description='ChatLearn Arguments',
                                     allow_abbrev=False)

    parser.add_argument("-c", "--config",
                        required=False,
                        help="where to load YAML configuration",
                        metavar="FILE")

    args, _ = parser.parse_known_args()

    if args.config:
        config_dir = os.path.dirname(args.config)
        args_yaml = parse_args_from_yaml(args.config, config_dir)
    else:
        config_dir = None
        args_yaml = None
    config = Config(args_yaml, config_dir)

    return config


class BaseConfig:
    """Base class includes some common format functions."""

    def __init__(self):
        self._finalize = True

    def __str__(self):
        members = [attr for attr in dir(self) \
                   if not callable(getattr(self, attr)) and not attr.startswith("__")]
        ser_str = self.__class__.__name__ + " {\n"
        for key in members:
            if key.startswith('_'):
                continue
            attr = getattr(self, key)
            attr = '"{}"'.format(attr) if isinstance(attr, str) else attr
            ser_str += "    %s = %s,\n" % (key, attr)
        ser_str += "}"

        return ser_str

    def __repr__(self):
        return self.__str__()

    def validate(self):
        pass


class SubConfig(BaseConfig):
    """Sub Config"""
    _is_changed = False

    def __setattr__(self, name, value):
        if not name.startswith("_") and getattr(self, name) != value:
            self._is_changed = True
        super().__setattr__(name, value)

    def is_changed(self):
        return self._is_changed


class LoraConfig(SubConfig):
    """Config for lora"""
    #: enable lora, default False.
    enable_lora: bool = False
    #: The "name_scope" parameter is used to specify a particular module to be converted to its LoRA.
    #: By default, it is set to None, which means there is no restriction on the module and any module
    #: can be converted using the "lora_layer" parameter. However, if "name_scope" is set to a specific
    #: value (e.g., "encoder"), only the modules whose name_scope contains the value "encoder" will be converted to LoRA.
    part_module_name: str = None
    #: The rank value of the LoRA, which is the r dimension of the A/B matrix.
    lora_dim: int = 8
    #: The LoRA dropout ratio refers to whether dropout computation is inserted in the forward pass
    #: of the LoRA layer. By default, the dropout ratio is set to 0.0.
    lora_dropout: float = 0.0
    #: When adding the values of the LoRA A and B matrices to the original weight matrix,
    #: the scaling value is set as "W = W + A * B * lora_scaling". By default, the scaling value
    #: is set to 1.0.
    lora_scaling: float = 1.0
    #: The layer class names involved in LoRA training in the model, separated by commas.
    lora_layer: str = LORA_LAYER
    #: LoRA training is enabled only in the ColumnParallelLinear layer of the MHA QKV module.
    column_only_qkv: bool = False


class BatchGenerationConfig(SubConfig):
    """Config for batch generation ranking and memory-efficiency."""

    #: [optional] sort prompts by length each episode.
    ranking: bool = False
    #: [optional] min prompt length in the first stage of batch generation.
    min_prompt_length: int = 0


class ModelConfig(BaseConfig):
    """Config for model."""

    #: [legacy] number of GPU used for one model, default 0.
    num_device: int = 0
    #: [required] number of GPU used for one model, default 0, same as num_device
    num_gpu: int = 0
    #: [required] number of GPU used for one model, default 0
    num_cpu: int = 0
    #: [optional] gpu per process, e.g., for PyTorch DDP, Megatron, DeepSpeed, `gpu_per_process` is set to 1
    gpu_per_process: int = None
    #: [optional] cpu per process
    cpu_per_process: int = None
    #: [optional] number of module replica,
    #: for gpu model, num_replica = num_gpu // (TP * PP * DP * EP),
    #: for cpu model, num_replica = num_cpu // cpu_per_process
    num_replica: int = 1
    #: [required] whether model is trainable
    trainable: bool = False
    #: [optional] tensor model parallel size
    tensor_model_parallel_size: int = None
    #: [optional] pipeline model parallel size
    pipeline_model_parallel_size: int = None
    #: [optional] expert model parallel size for Megatron-Core
    expert_model_parallel_size: int = None
    #: [optional] zero size
    zero_size: int = None
    #: [optional] FSDP parallel size
    fsdp_size: int = None
    #: [optional] config file for model
    model_config_file: str = ""
    config_dir: str = ""
    #: [optional] model type, e.g., Torch/Tensorflow, etc
    model_type: str = ""
    #: [optional] placeholder for other args
    args_dict: dict = None
    #: [optional] generation batch size, will overwrite generation batch size in RuntimeConfig
    generation_batch_size: int = None
    #: lora config
    lora: LoraConfig = None
    #: batch generation config
    batch_generation: BatchGenerationConfig = None
    #: offload optimizer states
    offload_optimizer_states = False
    #: parameter sync frequency
    sync_frequency = 1
    #: offload weights
    offload_weights = False
    #: free grad buffers
    free_grad_buffers = False
    #: overall switch for offload optimizer states/weights and free grad buffers
    free_memory = False
    #: force to free memory
    force_free_memory = False

    def __init__(self):
        super().__init__()
        self.args_dict = {}
        self.lora = LoraConfig()
        self.batch_generation = BatchGenerationConfig()

    def __str__(self):
        members = [attr for attr in dir(self) \
                   if not callable(getattr(self, attr)) and not attr.startswith("__")]
        ser_str = self.__class__.__name__ + " {\n"
        for key in members:
            if key.startswith('_'):
                continue
            attr = getattr(self, key)
            if key in ["lora", "batch_generation"]:
                if not attr.is_changed():
                    continue
            attr = '"{}"'.format(attr) if isinstance(attr, str) else attr
            ser_str += "    %s = %s,\n" % (key, attr)
        ser_str += "}"

        return ser_str


class RuntimeConfig(BaseConfig):
    """training related configs."""

    #: [required] number of episodes. One episode includes a inference and training loop.
    num_episode: int = 5000
    #: [required] number of samples per episode.
    sample_per_episode: int = 1000
    #: [optional] number of training epoch per episode. default set to 1.
    num_training_epoch: int = 1
    #: [optional] max iteration per sample, for mcts-style search algorithm
    max_iteration_per_sample = 1
    #: [required] generation(inference) batch size.
    generation_batch_size: int = 2
    #: [required] training micro batch size.
    train_micro_batch_size: int = 2
    #: [required] training global batch size.
    train_global_batch_size: int = None
    #: [required] save checkpoint per `save_episode_interval` episodes.
    save_episode_interval: int = None
    #: [optional] log time and memory per `log_interval` iterations.
    log_interval: int = 1
    #: [required]: data_path for dataset or a List of data_path for different kind of datasets
    data_path: Optional[Union[List[str], str]] = None
    #: [optional]: the ratio for each kind of data_path in a training episode, default: None
    data_ratio: Optional[Union[List[int], int]] = None
    #: [optional]: shuffle in each epoch of dataset, default: True
    data_shuffle: Optional[bool] = True
    #: [optional]: rerank batch of data by row, default: True
    data_rerank: Optional[bool] = False
    #: [optional]: colocate models into the same device
    colocation: List[str] = []
    #: [optional]: eval every N episode, if 0, will not eval
    eval_episode_interval: int = 0
    #: [optional]: enable resume training when data checkpoint is set
    enable_resume_training: bool = True
    #: [optional]: checkpoint for dataloader
    data_checkpoint_path: str = None
    #: [optional]: max data checkpoint nums
    max_data_ckpt_nums: int = None
    #: [optional]: load data checkpoint from iteration
    load_data_checkpoint_iteration: int = None
    #: [optional]: stream_data_loader type, ["fixed", "dynamic"]
    stream_data_loader_type: str = "fixed"
    #: private
    debug: bool = False
    #: enable nsys nvtx
    nsys: bool = False
    #: profiler dir
    profiler_dir: str = None
    #: coalesce_buffer size in mb
    coalesced_buffer_mb: int = 100
    #: concurrent parameter sync
    concurrent_comm: bool = True
    #: parameter sync communication type, broadcast/p2p
    param_sync_comm_type: str = PARAM_SYNC_COMM_TYPE.BROADCAST.value
    #: parameter sync max workers
    param_sync_max_workers: int = None
    #: communication type to regroup routed experts, allgather/alltoall
    routed_expert_regrouping_comm_type: str = ROUTED_EXPERT_REGROUPING_COMM_TYPE.ALLTOALL
    #: max number of relay episodes, if `max_relay_episode` is set to -1, then relay all episodes
    #: if `max_relay_episode` is set to 0, then relay is disabled
    max_relay_episode: int = 0
    #: relay after n episodes
    relay_episode_offset: int = 0
    #: training shuffle mode
    training_shuffle_mode: str = TrainingShffuleMode.BATCH
    #: consumed samples
    consumed_samples: int = 0
    #: concurrent model setup
    concurrent_setup: bool = False
    #: bucket size in the memory manager to reduce peak memory
    bucket_size_mb_in_memory_manager: int = 1024
    #: free collective group after parameter synchronization and rebuild before next synchronization
    free_sync_collective_group: bool = False
    #: [optional] cpu only model schedule policy, PACK or SPREAD
    #: PACK: All provided bundles are packed onto a single node on a best-effort basis.
    #: SPREAD: Each bundle is spread onto separate nodes on a best-effort basis.
    cpu_schedule_strategy: str = RAY_PG_STRATEGY.SPREAD.value
    #: exp name for each run
    exp_name: str = "CHATLEARN"
    #: output dir
    output_dir: str = "./"
    #: validate param sync
    validate_param_sync: bool = False
    #: whether to eval before training
    enable_eval_before_training: bool = False
    #: policy to regroup queue
    policy_to_regroup_queue: str = "global_barrier"
    #: configuration file path for logging
    log_config_file: str = ""
    #: [optional] placeholder for log_args_dict
    log_args_dict: dict = None

    def __init__(self):
        super().__init__()
        self._args_dict = {}

    def get(self, key):
        """
        Get other config by key.

        Args
        ----
        key: str
            key to get config
        """
        if key not in self._args_dict:
            logger.warning(f"{key} not found in RuntimeConfig")
        else:
            return self._args_dict[key]

    def validate(self):
        """
        :meta private:
        """
        for key in self._args_dict:
            if key == "save_interval":
                raise Exception("save_interval is deprecated, please use save_episode_interval to save checkpoints")


class RuntimeEnvConfig(BaseConfig):
    """Runtime env config, you can refer https://docs.ray.io/en/latest/ray-core/handling-dependencies.html for more information."""

    #: pip install packages
    pip: List[str] = []
    #: python modules
    py_modules: List[str] = []
    #: working directory
    working_dir: str = os.getcwd()
    #: platform, e.g., DLC
    platform: str = ""
    #: excludes files from packaging
    excludes: List[str] = []

    def __init__(self):
        super().__init__()
        self._args_dict = {}

    def get(self, key):
        """
        Get other config by key

        Args
        ----
        key: str
            Key to get config.
        """
        if key not in self._args_dict:
            logger.warning(f"{key} not found in RuntimeConfig")
        else:
            return self._args_dict[key]


class Config(BaseConfig):
    """A class to manage chatlearn configuration.

    Args
    ----
      param_dict: dict
      dict format of parameters."""

    def __init__(self, param_dict=None, config_dir=None):
        super().__init__()
        self._finalize = False
        self.models = {}
        self.env_args = RuntimeEnvConfig()
        self.runtime_args = RuntimeConfig()
        self.config_dir = config_dir
        self._active_module_args = None

        self.initialized = False
        if param_dict:
            self._parse_params(param_dict)
            self._validate_params()
        # remove later, just for compatibility
        self.rlhf_args = self.runtime_args
        self._finalize = True

    def _parse_params(self, param_dict):
        """Parse params from param_dict."""

        def set_param(user_args, config_cls, instance):
            for attribute, default_value in get_attributes(config_cls):
                if attribute in user_args:
                    value = user_args[attribute]
                    if attribute == "colocation":
                        colocation_list = []
                        for group in value:
                            colocation_list.append(group.replace(' ', '').split(','))
                        value = colocation_list
                    elif attribute == "data_ratio":
                        if isinstance(value, str):
                            value = [int(v) for v in value.split(',')]
                else:
                    value = default_value
                original_value = getattr(instance, attribute)
                if original_value is not None:
                    assert isinstance(original_value, type(value)), \
                        f"{instance}.{attribute} should be type of {type(original_value)} but got {type(value)}"

                setattr(instance, attribute, value)
            for user_attribute in user_args:
                if not hasattr(config_cls, user_attribute):
                    if hasattr(instance, "_args_dict"):
                        getattr(instance, "_args_dict")[user_attribute] = user_args[user_attribute]
                    else:
                        raise RuntimeError(f"attribute {user_attribute} not defined in {config_cls.__name__}")
            instance.validate()

        for model_name, model_args in param_dict["models"].items():
            model_config = ModelConfig()
            model_config.config_dir = self.config_dir
            for user_attribute, user_value in model_args.items():
                if hasattr(ModelConfig, user_attribute):
                    original_value = getattr(ModelConfig, user_attribute)
                    if 'num_device' == user_attribute:
                        logger.warning("num_device is deprecated, please use num_gpu instead")
                        if 'num_gpu' not in model_args.keys():
                            setattr(model_config, "num_gpu", user_value)
                        else:
                            logger.warning("both num_device and num_gpu are set, use num_gpu")
                            continue
                    if 'lora' == user_attribute:
                        set_param(user_value, LoraConfig, model_config.lora)
                        user_value = model_config.lora
                    elif "batch_generation" == user_attribute:
                        set_param(user_value, BatchGenerationConfig, model_config.batch_generation)
                        user_value = model_config.batch_generation
                    if original_value is not None:
                        assert isinstance(user_value, type(original_value)), \
                            f"ModelConfig.{user_attribute} should be type of {type(original_value)} but got {type(user_value)} ({user_value})"
                    setattr(model_config, user_attribute, user_value)
                else:
                    logger.warning(f"unknown argument {user_attribute}")

            self.models[model_name] = model_config
            if model_config.model_config_file:
                model_config.model_config_file = get_path(model_config.model_config_file, self.config_dir)
                model_config.args_dict = parse_args_from_yaml(model_config.model_config_file, self.config_dir)
        if "runtime" in param_dict:
            set_param(param_dict["runtime"], RuntimeConfig, self.runtime_args)
        elif "rlhf" in param_dict:
            logger.warning("rlhf is deprecated, please use runtime as section name")
            set_param(param_dict["rlhf"], RuntimeConfig, self.runtime_args)
        if "runtime_env" in param_dict:
            set_param(param_dict["runtime_env"], RuntimeEnvConfig, self.env_args)

        if self.runtime_args.log_config_file:
            self.runtime_args.log_config_file = get_path(self.runtime_args.log_config_file, self.config_dir)
            self.runtime_args.log_args_dict = parse_args_from_yaml(self.runtime_args.log_config_file, self.config_dir)

        def _get_and_check_type(value, default_value, key):
            # To be noticed: all str type values should in lower case.
            if isinstance(value, str):
                value = value.lower()
            if default_value is None:
                return value
            if not isinstance(value, type(default_value)):
                raise ValueError("%s type error, expected: %s." \
                                 % (key, type(default_value)))
            return value

    def _validate_params(self):
        if self.runtime_args.train_global_batch_size is None:
            self.runtime_args.train_global_batch_size = self.runtime_args.train_micro_batch_size
        assert self.runtime_args.train_global_batch_size % self.runtime_args.train_micro_batch_size == 0, \
            f"train_global_batch_size should be times of train_micro_batch_size," \
            f"but got {self.runtime_args.train_global_batch_size}/{self.runtime_args.train_micro_batch_size}"
        assert self.runtime_args.train_global_batch_size <= self.runtime_args.sample_per_episode, \
            "train_global_batch_size should be less than or equal to sample_per_episode, " \
            f"got {self.runtime_args.train_global_batch_size} and {self.runtime_args.sample_per_episode}"
        assert self.runtime_args.stream_data_loader_type.lower() in ["fixed", "dynamic"]
        assert self.runtime_args.cpu_schedule_strategy in [strategy.value for strategy in RAY_PG_STRATEGY]
        assert self.runtime_args.param_sync_comm_type in list(PARAM_SYNC_COMM_TYPE)
        if isinstance(self.runtime_args.data_path, list):
            assert self.runtime_args.data_ratio is not None and isinstance(self.runtime_args.data_ratio, list), (
                f"expect data_ratio to be list when data_path is list, got {self.runtime_args.data_ratio}"
            )
            assert len(self.runtime_args.data_path) == len(self.runtime_args.data_ratio), (
                "expect data_path and data_ratio to have same length, "
                f"got {len(self.runtime_args.data_path)} and {len(self.runtime_args.data_ratio)}"
            )
        for model_name, model_args in self.models.items():
            if model_args.num_gpu >= 1:
                if model_args.gpu_per_process is None:
                    model_args.gpu_per_process = 1
                else:
                    assert model_args.gpu_per_process <= model_args.num_gpu, \
                        f"{model_name}: gpu_per_process: {model_args.gpu_per_process}, num_cpu: {model_args.num_gpu}"
            elif model_args.num_cpu >= 1:
                if model_args.cpu_per_process is None:
                    model_args.cpu_per_process = 1
                else:
                    assert model_args.cpu_per_process <= model_args.num_cpu, \
                        f"{model_name}: cpu_per_process: {model_args.cpu_per_process}, num_cpu: {model_args.num_cpu}"
            if model_args.generation_batch_size is not None and model_args.generation_batch_size <= 0:
                model_args.generation_batch_size = DYNAMIC_BATCH_SIZE
            if model_args.generation_batch_size is None:
                if self.runtime_args.generation_batch_size:
                    model_args.generation_batch_size = self.runtime_args.generation_batch_size
            for key in ["pipeline_model_parallel_size", "tensor_model_parallel_size", "zero_size"]:
                if model_args.args_dict.get(key) is not None:
                    setattr(model_args, key, model_args.args_dict.get(key))
                    assert getattr(model_args, key) >= 1
                elif getattr(model_args, key) is None:
                    setattr(model_args, key, 1)

            for key in ["fsdp_size"]:
                if getattr(model_args, key) is not None:
                    setattr(model_args, key, getattr(model_args, key))
                    if getattr(model_args, key) == -1:
                        print(f"set_fsdp_size {getattr(model_args, key)} to num_gpu: {model_args.num_gpu}")
                        setattr(model_args, key, model_args.num_gpu)
                    assert getattr(model_args, key) >= 1
                elif getattr(model_args, key) is None:
                    setattr(model_args, key, 1)

            ep_size = model_args.args_dict.get("expert_model_parallel_size")
            moe_ep_size = model_args.args_dict.get("moe_expert_model_parallel_size")
            if ep_size is not None and moe_ep_size is not None:
                assert ep_size == moe_ep_size, (
                    f"{model_name}: if you set moe_expert_model_parallel_size ({moe_ep_size}), "
                    f"it must be equal to expert_model_parallel_size ({ep_size})"
                )
                finalized_ep_size = ep_size
            elif ep_size is not None:
                finalized_ep_size = ep_size
            elif moe_ep_size is not None:
                finalized_ep_size = moe_ep_size
            else:
                finalized_ep_size = 1
            assert finalized_ep_size >= 1
            setattr(model_args, "expert_model_parallel_size", finalized_ep_size)

            if model_args.tensor_model_parallel_size > 1 or model_args.pipeline_model_parallel_size > 1 or model_args.expert_model_parallel_size > 1:
                assert model_args.zero_size == 1 or model_args.zero_size is None
                assert model_args.fsdp_size == 1 or model_args.fsdp_size is None
                assert model_args.num_gpu % (
                    model_args.tensor_model_parallel_size * model_args.pipeline_model_parallel_size * model_args.expert_model_parallel_size) == 0, \
                    f"{model_name}: num_gpu must be divisible by tensor_model_parallel_size * pipeline_model_parallel_size * " \
                    f"expert_model_parallel_size, but got num_gpu = {model_args.num_gpu}, " \
                    f"tensor_model_parallel_size = {model_args.tensor_model_parallel_size}, " \
                    f"pipeline_model_parallel_size = {model_args.pipeline_model_parallel_size}, and "\
                    f"expert_model_parallel_size = {model_args.expert_model_parallel_size}."
            assert model_args.num_gpu > 0 or model_args.num_cpu > 0, \
                f"{model_name} num_gpu: {model_args.num_gpu}, num_cpu: {model_args.num_cpu}, at least one of them should be set"

            if model_args.num_gpu >= 1:
                if model_args.zero_size > 1:
                    assert model_args.num_gpu % model_args.zero_size == 0
                    model_args.num_replica = model_args.num_gpu // model_args.zero_size
                elif model_args.fsdp_size > 1:
                    # For FSDP, num_gpu must be divisible by fsdp_size
                    assert model_args.num_gpu % model_args.fsdp_size == 0
                    model_args.num_replica = model_args.num_gpu // (
                        model_args.tensor_model_parallel_size * model_args.pipeline_model_parallel_size \
                            * model_args.expert_model_parallel_size * model_args.fsdp_size)
                else:
                    model_args.num_replica = model_args.num_gpu // (
                        model_args.tensor_model_parallel_size * model_args.pipeline_model_parallel_size * model_args.expert_model_parallel_size)

            elif model_args.num_cpu >= 1:
                model_args.num_replica = model_args.num_cpu // model_args.cpu_per_process
            assert model_args.num_replica * model_args.generation_batch_size <= self.runtime_args.sample_per_episode, \
                f"{model_name}: num_replica * batch_size {model_args.num_replica}*{model_args.generation_batch_size} " + \
                f"should be less than or equal to sample_per_episode {self.runtime_args.sample_per_episode}"
            if model_args.batch_generation.min_prompt_length:
                logger.info(f"Enable batch generation: \
                    min_prompt_length = {model_args.batch_generation.min_prompt_length}")
            if model_args.free_memory:
                model_args.offload_weights = True
                if model_args.trainable:
                    model_args.free_grad_buffers = True
                    model_args.offload_optimizer_states = True
        if self.runtime_args.colocation and len(self.runtime_args.colocation) > 0:
            model_set = set()
            for colocate_models in self.runtime_args.colocation:
                for model_name in colocate_models:
                    assert model_name not in model_set, f"Model {model_name} should only appear once in colocation group"
                    model_set.add(model_name)
        if self.runtime_args.exp_name not in self.runtime_args.output_dir:
            self.runtime_args.output_dir = f"{self.runtime_args.output_dir}/{self.runtime_args.exp_name}"
        logger.info(f"Env Config: \n{self.env_args}")
        logger.info(f"Runtime Config: \n{self.runtime_args}")
        for name, model_args in self.models.items():
            logger.info(f"Model({name}) Config: \n{model_args}")

    @property
    def active_module_args(self):
        return self._active_module_args

    @active_module_args.setter
    def active_module_args(self, config):
        self._active_module_args = config