diff --git a/Dockerfile b/Dockerfile index d8e375da7..bf6517852 100644 --- a/Dockerfile +++ b/Dockerfile @@ -108,8 +108,11 @@ COPY . . RUN uv pip install grpcio-tools RUN cd /app && find tunix/experimental/distributed -name "*.proto" -exec python -m grpc_tools.protoc -I/app --python_out=/app --grpc_python_out=/app {} + -# Install Tunix in editable mode -RUN uv pip install --no-deps -e . +# Install Tunix in editable mode and register MaxText vLLM adapter + Raiden +RUN uv pip install --no-deps -e . && \ + RAIDEN_WHEEL_DIR="/app/raiden_wheels" bash /app/scripts/install_raiden.sh && \ + uv pip install --no-deps /opt/venv/lib/python3.12/site-packages/maxtext/integration/vllm && \ + uv pip install git+https://github.com/mlcommons/logging.git diff --git a/tunix/experimental/examples/math_gsm8k_dist/k8s_launcher.sh b/tunix/experimental/examples/math_gsm8k_dist/k8s_launcher.sh index ba494e87a..89a95578f 100755 --- a/tunix/experimental/examples/math_gsm8k_dist/k8s_launcher.sh +++ b/tunix/experimental/examples/math_gsm8k_dist/k8s_launcher.sh @@ -107,6 +107,9 @@ export WANDB_RUN_NAME=${WANDB_RUN_NAME:-} export WANDB_API_KEY=${WANDB_API_KEY:-} export LOG_DIR=${LOG_DIR:-} export TRAJECTORY_LOG_DIR=${TRAJECTORY_LOG_DIR:-} +export RCP_LOGGING=${RCP_LOGGING:-false} +export METRIC_LOGGER_DIR=${METRIC_LOGGER_DIR:-} +export TARGET_ACCURACY=${TARGET_ACCURACY:-0.69} export TFDS_DATA_DIR=${TFDS_DATA_DIR:-"artifacts/data"} export TFDS_SPLIT=${TFDS_SPLIT:-train} export FLUSH_METRICS_EVERY_N_STEPS=${FLUSH_METRICS_EVERY_N_STEPS:-1} @@ -180,6 +183,10 @@ start_orchestrator() { if [[ "${DEBUG}" == "1" || "${DEBUG}" == "true" || "${DEBUG}" == "True" ]]; then debug_flag="--debug" fi + local rcp_flag="" + if [[ "${RCP_LOGGING}" == "1" || "${RCP_LOGGING}" == "true" || "${RCP_LOGGING}" == "True" ]]; then + rcp_flag="--rcp_logging" + fi "$PYTHON" "$YAML_GEN" \ "$YAML_DIR/jobset.cpu.yaml" \ @@ -221,6 +228,20 @@ start_orchestrator() { ${MAX_SEQ_TOKEN_PER_TPU:+--max_seq_token_per_tpu=${MAX_SEQ_TOKEN_PER_TPU}} \ ${MAX_SEGMENTS_PER_PACKED_ROW:+--max_segments_per_packed_row=${MAX_SEGMENTS_PER_PACKED_ROW}} \ ${TRAINER_MESH_FSDP:+--trainer_fsdp=${TRAINER_MESH_FSDP}} \ + --eval_every_n_steps=${EVAL_EVERY_N_STEPS} \ + --learning_rate=${LEARNING_RATE} \ + --b1=${ADAM_B1} \ + --b2=${ADAM_B2} \ + --weight_decay=${WEIGHT_DECAY} \ + --max_grad_norm=${MAX_GRAD_NORM} \ + --train_mesh_tp=${TRAINER_MESH_TP} \ + --train_mesh_expert=${TRAINER_MESH_EXPERT} \ + --rollout_mesh_tp=${ROLLOUT_MESH_TP} \ + --rollout_engine=${SAMPLER} \ + --tpu_topology="${TRAINER_TPU_SLICE}+${ROLLOUT_TPU_SLICE}" \ + --target_accuracy=${TARGET_ACCURACY} \ + ${METRIC_LOGGER_DIR:+--metric_logger_dir="${METRIC_LOGGER_DIR}"} \ + ${rcp_flag} \ ${debug_flag} \ " \ | apply_manifest diff --git a/tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py b/tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py index 82e537ba1..6428e1709 100644 --- a/tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py +++ b/tunix/experimental/examples/math_gsm8k_dist/run_gsm8k_dist_grpo.py @@ -57,6 +57,7 @@ from tunix.experimental.worker import remote_execution # pylint: disable=g-import-not-at-top from tunix.rl import algorithm_config # pylint: disable=g-import-not-at-top from tunix.sft import metrics_logger as metrics_logger_lib # pylint: disable=g-import-not-at-top +from tunix.utils import mllog_utils # pylint: disable=g-import-not-at-top ProcessContext = runtime_context.ProcessContext @@ -288,6 +289,79 @@ def _parse_args(argv: list[str]) -> argparse.Namespace: action="store_true", help="Enable debug logging and print full sampler responses.", ) + parser.add_argument( + "--rcp_logging", + action="store_true", + default=False, + help="Enable MLPerf RCP (mllog) compliance logging.", + ) + parser.add_argument( + "--metric_logger_dir", + type=str, + default=os.getenv("METRIC_LOGGER_DIR", None), + help="Directory or GCS URI for MLPerf RCP output (seed_.out).", + ) + parser.add_argument( + "--target_accuracy", + type=float, + default=float(os.getenv("TARGET_ACCURACY", "0.69")), + help="Target evaluation accuracy for MLPerf RCP compliance logging.", + ) + parser.add_argument( + "--eval_every_n_steps", + type=int, + default=int(os.getenv("EVAL_EVERY_N_STEPS", "1000000")), + ) + parser.add_argument( + "--learning_rate", + type=float, + default=float(os.getenv("LEARNING_RATE", "2.0e-7")), + ) + parser.add_argument( + "--b1", + type=float, + default=float(os.getenv("ADAM_B1", "0.9")), + ) + parser.add_argument( + "--b2", + type=float, + default=float(os.getenv("ADAM_B2", "0.999")), + ) + parser.add_argument( + "--weight_decay", + type=float, + default=float(os.getenv("WEIGHT_DECAY", "0.01")), + ) + parser.add_argument( + "--max_grad_norm", + type=float, + default=float(os.getenv("MAX_GRAD_NORM", "1.0")), + ) + parser.add_argument( + "--train_mesh_tp", + type=int, + default=int(os.getenv("TRAINER_MESH_TP", "1")), + ) + parser.add_argument( + "--train_mesh_expert", + type=int, + default=int(os.getenv("TRAINER_MESH_EXPERT", "1")), + ) + parser.add_argument( + "--rollout_mesh_tp", + type=int, + default=int(os.getenv("ROLLOUT_MESH_TP", "1")), + ) + parser.add_argument( + "--rollout_engine", + type=str, + default=os.getenv("SAMPLER", "vllm"), + ) + parser.add_argument( + "--tpu_topology", + type=str, + default=os.getenv("TPU_TOPOLOGY", None), + ) return parser.parse_args(argv) @@ -389,6 +463,8 @@ def main(argv: list[str], context: ProcessContext | None = None) -> None: ) args = _parse_args(argv) + if args.rcp_logging: + mllog_utils.init_start(args) if args.debug: # Enable canonical debug logging and print full sampler responses logging.getLogger().setLevel(logging.DEBUG) @@ -520,16 +596,39 @@ def main(argv: list[str], context: ProcessContext | None = None) -> None: step, step, ), - on_step_end=lambda step, result: logging.info( - "<<< Step %d finished | Advanced to Policy Version: %d", - step, - step + 1, + on_step_end=lambda step, result: ( + logging.info( + "<<< Step %d finished | Advanced to Policy Version: %d", + step, + step + 1, + ), + mllog_utils.log_rcp_step_stats( + program.metrics_logger, + args=args, + step=step + 1, + ) + if args.rcp_logging + else None, ), ) + if args.rcp_logging: + train_ds = gsm8k.load_gsm8k_dataset( + split=args.tfds_split, + data_dir=args.tfds_data_dir, + shuffle=args.shuffle, + seed=args.seed, + ) + mllog_utils.init_print( + args, + train_dataset=train_ds, + ) + try: logging.info("Bringing up remote workers through ClusterOrchestrator...") cluster.bring_up_workers(dummy_data=None) + if args.rcp_logging: + mllog_utils.train_start(args, step=0) logging.info( "Cluster workers ready: %s. Starting StandardRLProgram execution...", [w.worker_id for w in cluster.worker_infos()], @@ -539,7 +638,21 @@ def main(argv: list[str], context: ProcessContext | None = None) -> None: num_steps=args.max_steps, bring_up=False, ) + if args.rcp_logging: + completed_steps = ( + program.last_step_result.step + 1 + if program.last_step_result is not None + else args.max_steps + ) + mllog_utils.train_stop(args, step=completed_steps, status="success") except BaseException as exc: + if args.rcp_logging: + completed_steps = ( + program.last_step_result.step + 1 + if program.last_step_result is not None + else 0 + ) + mllog_utils.train_stop(args, step=completed_steps, status="aborted") logging.exception("FATAL: StandardRLProgram execution failed: %s", exc) raise finally: diff --git a/tunix/experimental/orchestrator/rl_program.py b/tunix/experimental/orchestrator/rl_program.py index 5311f0769..4a2019ef4 100644 --- a/tunix/experimental/orchestrator/rl_program.py +++ b/tunix/experimental/orchestrator/rl_program.py @@ -625,6 +625,9 @@ def _collect_and_log_step_metrics( consumed_policy_version: int, log_step: int, sampler_agreement: dict[str, tuple[Any, list[float]]] | None = None, + policy_training_time: float = 0.0, + exposed_generation_time: float = 0.0, + weight_sync_time: float = 0.0, ) -> dict[str, Any]: """Logs rollout, reward, trainer, and orchestrator metrics. @@ -744,6 +747,29 @@ def _collect_and_log_step_metrics( self.mode, log_step, ) + self.metrics_logger.log( + self.metrics_prefix, + "rollout/global_valid_toks", + float(np.sum(total_lengths)), + self.mode, + log_step, + ) + elif completion_lengths: + self.metrics_logger.log( + self.metrics_prefix, + "rollout/global_valid_toks", + float(np.sum(completion_lengths)), + self.mode, + log_step, + ) + if all_step_items: + self.metrics_logger.log( + self.metrics_prefix, + "rollout/global_valid_seqs", + float(len(all_step_items)), + self.mode, + log_step, + ) if turns_list: self.metrics_logger.log( self.metrics_prefix, @@ -839,6 +865,9 @@ def _collect_and_log_step_metrics( "num_rollouts": float(num_rollouts), "num_microbatches": float(num_microbatches), "step_time_sec": float(step_time_sec), + "policy_training_time": float(policy_training_time), + "exposed_generation_time": float(exposed_generation_time), + "weight_sync_time": float(weight_sync_time), } for tag, val in orchestrator_stats.items(): self.metrics_logger.log( @@ -1101,6 +1130,9 @@ async def train_stage(self) -> None: groups_consumed = 0 checkpoint_saved = False final_minibatch_completed = False + policy_training_time = 0.0 + exposed_generation_time = 0.0 + weight_sync_time = 0.0 async def _maybe_save_checkpoint() -> None: nonlocal checkpoint_saved @@ -1123,7 +1155,9 @@ async def _maybe_save_checkpoint() -> None: checkpoint_saved = True while groups_consumed < self.full_batch_size: + _t_gen = time.monotonic() scored_items = await self.scored_q.get_batch(num_groups=1) + exposed_generation_time += time.monotonic() - _t_gen if not scored_items: assembled_batches = self.assembler.flush() else: @@ -1184,16 +1218,20 @@ async def _maybe_save_checkpoint() -> None: len(mb.trajectory_ids), logging_utils.summarize_list(list(mb.trajectory_ids)), ) + _t_train = time.monotonic() step_result = await self.engine.train_step( batch, role=datatypes.Role.ACTOR, accumulate_gradients=True, apply_optimizer=mb.is_final_batch, ) + policy_training_time += time.monotonic() - _t_train if mb.is_final_batch: + _t_metrics = time.monotonic() trainer_metrics = await self.engine.get_metrics( role=datatypes.Role.ACTOR ) + policy_training_time += time.monotonic() - _t_metrics final_minibatch_completed = True # TODO(tunix-dev): Configurable checkpointing frequency. Today we # checkpoint at the same frequency as the weight update. @@ -1221,7 +1259,9 @@ async def _maybe_save_checkpoint() -> None: break if self.sync_weights: + _t_sync = time.monotonic() new_version = await self.engine.sync_weights(role=datatypes.Role.ACTOR) + weight_sync_time = time.monotonic() - _t_sync self.policy_version = ( new_version if new_version is not None else self.policy_version + 1 ) @@ -1251,6 +1291,9 @@ async def _maybe_save_checkpoint() -> None: consumed_policy_version=consumed_policy_version, log_step=current_step, sampler_agreement=step_sampler_agreement, + policy_training_time=policy_training_time, + exposed_generation_time=exposed_generation_time, + weight_sync_time=weight_sync_time, ) self._log_consumed_trajectories( all_step_items, diff --git a/tunix/utils/mllog_utils.py b/tunix/utils/mllog_utils.py index a447d827b..238c08bc0 100644 --- a/tunix/utils/mllog_utils.py +++ b/tunix/utils/mllog_utils.py @@ -30,24 +30,93 @@ mllogger = None +_gcs_target_path: Optional[str] = None +_local_log_path: Optional[str] = None + + +def _parse_topology_devices( + topology: Optional[str], rollout_replicas: int = 1 +) -> Optional[int]: + """Parses TPU slice strings like 'tpuv5p:2x2x2+tpuv5p:2x2x1' into total chip count.""" + if not topology: + return None + parts = str(topology).split("+") + total = 0 + for idx, part in enumerate(parts): + dims_str = part.split(":")[-1] + dims = [int(x) for x in dims_str.split("x") if x.isdigit()] + if not dims: + continue + chips = 1 + for d in dims: + chips *= d + if idx > 0: + chips *= max(1, int(rollout_replicas)) + total += chips + return total if total > 0 else None + + +def _flush_to_gcs_if_needed() -> None: + """Copies the local mllog file to GCS if metric_logger_dir was a gs:// URI.""" + if not (_is_master_process() and _gcs_target_path and _local_log_path): + return + if not os.path.exists(_local_log_path): + return + try: + for h in getattr(mllogger.logger, "handlers", []): + if isinstance(h, logging.FileHandler): + h.flush() + try: + import fsspec # pylint: disable=g-import-not-at-top + + fs = fsspec.filesystem("gs") + fs.put(_local_log_path, _gcs_target_path) + except Exception: # pylint: disable=broad-exception-caught + import tensorflow as tf # pylint: disable=g-import-not-at-top + + tf.io.gfile.makedirs(os.path.dirname(_gcs_target_path)) + tf.io.gfile.copy(_local_log_path, _gcs_target_path, overwrite=True) + except Exception as exc: # pylint: disable=broad-exception-caught + logging.warning( + "Failed to copy mllog file %s to %s: %s", + _local_log_path, + _gcs_target_path, + exc, + ) + + def configure_logger( metric_logger_dir: Optional[str] = None, seed: Optional[int] = None, filename: Optional[str] = None, ): """Configures mllog output file if metric_logger_dir or filename is provided.""" + global _gcs_target_path, _local_log_path if not (_is_master_process() and mllog is not None and mllogger is not None): return + seed_val = seed if seed is not None else 1 if filename is None and metric_logger_dir is not None: - if metric_logger_dir.endswith(".out") or metric_logger_dir.endswith(".log"): + if metric_logger_dir.startswith("gs://"): + if metric_logger_dir.endswith(".out") or metric_logger_dir.endswith( + ".log" + ): + _gcs_target_path = metric_logger_dir + else: + _gcs_target_path = os.path.join( + metric_logger_dir.rstrip("/"), f"seed_{seed_val}.out" + ) + filename = os.path.join("/tmp/rcp_logs", f"seed_{seed_val}.out") + elif metric_logger_dir.endswith(".out") or metric_logger_dir.endswith( + ".log" + ): filename = metric_logger_dir else: - seed_val = seed if seed is not None else 1 filename = os.path.join(metric_logger_dir, f"seed_{seed_val}.out") if filename is not None: abs_filename = os.path.abspath(filename) + _local_log_path = abs_filename os.makedirs(os.path.dirname(abs_filename), exist_ok=True) existing_files = [ os.path.abspath(getattr(h, "baseFilename", "")) @@ -132,8 +201,11 @@ def block_start(args=None, step: int = 0, samples_count: Optional[int] = None): if _is_master_process() and mllogger is not None: if samples_count is None and args is not None: global_batch_size = getattr(args, "batch_size", 1) * getattr(args, "num_generations", 1) - eval_interval = getattr(args, "eval_every_n_steps", getattr(args, "max_steps", 1)) - samples_count = eval_interval * global_batch_size + max_steps = getattr(args, "max_steps", None) + eval_interval = getattr(args, "eval_every_n_steps", max_steps if max_steps is not None else 1) + if max_steps is not None: + eval_interval = min(int(eval_interval), max(0, int(max_steps) - int(step))) + samples_count = int(eval_interval) * global_batch_size metadata = {"step": int(step)} if samples_count is not None: @@ -150,6 +222,7 @@ def train_start(args=None, step: int = 0, samples_count: Optional[int] = None): init_stop() run_start() block_start(args=args, step=step, samples_count=samples_count) + _flush_to_gcs_if_needed() def block_stop(step: int = 0, samples_count: Optional[int] = None): @@ -474,9 +547,117 @@ def _extract_kv_from_metrics_buffer(metrics_buffer: Any) -> dict[str, Any]: if val is not None: kv_stats[k] = val + # Case 4: tunix.sft.metrics_logger.MetricsLogger or StandardRLProgram + raw_metrics = getattr(metrics_buffer, "_metrics", None) + if raw_metrics is None and hasattr(metrics_buffer, "metrics_logger"): + raw_metrics = getattr(metrics_buffer.metrics_logger, "_metrics", None) + if isinstance(raw_metrics, dict): + for prefix_dict in raw_metrics.values(): + if not isinstance(prefix_dict, dict): + continue + for mode_key, mode_dict in prefix_dict.items(): + if str(mode_key) != "train" or not isinstance(mode_dict, dict): + continue + for k, vals in mode_dict.items(): + if vals and k not in kv_stats: + val = _clean_metric_val(vals[-1]) + if val is not None: + kv_stats[k] = val + if k.startswith("trainer/"): + short_k = k[len("trainer/") :] + if short_k not in kv_stats: + kv_stats[short_k] = val + return kv_stats +def log_rcp_step_stats( + metrics_source: Any, + args: Any = None, + step: int = 1, + samples_count: Optional[int] = None, + total_devices: Optional[int] = None, +) -> None: + """Emits the two MLPerf RCP tracked_stats events (train + timing) matching MLCommons seed_1.out.""" + if not (_is_master_process() and mllog is not None and mllogger is not None): + return + + stats = _extract_kv_from_metrics_buffer(metrics_source) + if not stats: + return + + step_num = int(step) + gbs = None + if args is not None: + gbs = getattr(args, "batch_size", 1) * getattr(args, "num_generations", 1) + if samples_count is None and gbs is not None: + samples_count = step_num * gbs + + # 1. Train stats event: reduced_train_loss, reward, grad_norm, global_valid_toks, global_valid_seqs + loss_val = stats.get("reduced_train_loss", stats.get("loss", stats.get("trainer/loss"))) + reward_val = stats.get( + "reward", + stats.get("rewards/mean", stats.get("trajectory_rewards/mean", stats.get("train_reward"))), + ) + grad_norm_val = stats.get("grad_norm", stats.get("trainer/grad_norm")) + valid_seqs = stats.get( + "global_valid_seqs", + stats.get("rollout/global_valid_seqs", stats.get("orchestrator/num_rollouts", float(gbs) if gbs else None)), + ) + valid_toks = stats.get("global_valid_toks", stats.get("rollout/global_valid_toks")) + if valid_toks is None and valid_seqs is not None: + mean_toks = stats.get( + "rollout/total_tokens_mean", + stats.get("rollout/completion_length_mean", stats.get("generation/completions/mean_raw_length")), + ) + if mean_toks is not None: + valid_toks = float(mean_toks) * float(valid_seqs) + + train_tracked = { + "reduced_train_loss": loss_val, + "reward": reward_val, + "grad_norm": grad_norm_val, + "global_valid_toks": float(valid_toks) if valid_toks is not None else None, + "global_valid_seqs": float(valid_seqs) if valid_seqs is not None else None, + } + log_tracked_stats(train_tracked, step=step_num, samples_count=samples_count) + + # 2. Timing stats event: train_step_time, policy_training_time, exposed_generation_time, weight_sync_time, valid_tokens_per_sec_per_gpu + step_time = stats.get( + "train_step_time", + stats.get("orchestrator/step_time_sec", stats.get("perf/global_step_time", stats.get("step_time"))), + ) + policy_time = stats.get("policy_training_time", stats.get("orchestrator/policy_training_time")) + exposed_gen_time = stats.get("exposed_generation_time", stats.get("orchestrator/exposed_generation_time")) + weight_sync_time = stats.get("weight_sync_time", stats.get("orchestrator/weight_sync_time")) + + if total_devices is None and args is not None: + total_devices = _parse_topology_devices( + getattr(args, "tpu_topology", None), + getattr(args, "rollout_replicas", 1), + ) + + toks_per_sec_per_gpu = stats.get("valid_tokens_per_sec_per_gpu") + if ( + toks_per_sec_per_gpu is None + and valid_toks is not None + and step_time is not None + and float(step_time) > 0 + and total_devices + ): + toks_per_sec_per_gpu = float(valid_toks) / (float(step_time) * float(total_devices)) + + timing_tracked = { + "train_step_time": step_time, + "policy_training_time": policy_time, + "exposed_generation_time": exposed_gen_time, + "weight_sync_time": weight_sync_time, + "valid_tokens_per_sec_per_gpu": toks_per_sec_per_gpu, + } + log_tracked_stats(timing_tracked, step=step_num, samples_count=samples_count) + _flush_to_gcs_if_needed() + + MLPERF_TRACKED_KEYS = frozenset({ "step_time", "train_reward", @@ -650,6 +831,7 @@ def run_stop(status: str = "success", samples_count: Optional[int] = None): key=getattr(constants, "RUN_STOP", "run_stop"), metadata=metadata, ) + _flush_to_gcs_if_needed() def init_print( @@ -671,14 +853,14 @@ def init_print( ) # Extract batch & step configs - batch_size = getattr(args, "batch_size", 8) - num_generations = getattr(args, "num_generations", 8) + batch_size = getattr(args, "batch_size", None) or 8 + num_generations = getattr(args, "num_generations", None) or 8 global_batch_size = batch_size * num_generations - mini_batch_size = getattr(args, "mini_batch_size", batch_size) - train_micro_batch_size = getattr(args, "train_micro_batch_size", 1) - max_steps = getattr(args, "max_steps", 50) - max_prompt_length = getattr(args, "max_prompt_length", 4096) - max_response_length = getattr(args, "max_response_length", 8192) + mini_batch_size = getattr(args, "mini_batch_size", None) or batch_size + train_micro_batch_size = getattr(args, "train_micro_batch_size", None) or 1 + max_steps = getattr(args, "max_steps", None) or 50 + max_prompt_length = getattr(args, "max_prompt_length", None) or 4096 + max_response_length = getattr(args, "max_response_length", None) or 8192 max_seq_len = max_prompt_length + max_response_length # Train / Eval sample counts @@ -767,17 +949,17 @@ def init_print( getattr(constants, "OPT_ADAMW_EPSILON", "opt_adamw_epsilon"): 1e-8, getattr(constants, "OPT_ADAMW_WEIGHT_DECAY", "opt_adamw_weight_decay"): getattr(args, "weight_decay", 0.01), getattr(constants, "OPT_GRADIENT_CLIP_NORM", "opt_gradient_clip_norm"): getattr(args, "max_grad_norm", 1.0), - getattr(constants, "OPT_LR_WARMUP_STEPS", "opt_learning_rate_warmup_steps"): 0, - getattr(constants, "OPT_LR_DECAY_STEPS", "opt_learning_rate_decay_steps"): max_steps, - getattr(constants, "OPT_LR_DECAY_SCHEDULE", "opt_learning_rate_decay_schedule"): "constant", + getattr(constants, "OPT_LR_WARMUP_STEPS", "opt_learning_rate_warmup_steps"): getattr(args, "warmup_steps", 0), + getattr(constants, "OPT_LR_DECAY_STEPS", "opt_learning_rate_decay_steps"): getattr(args, "lr_decay_steps", max_steps), + getattr(constants, "OPT_LR_DECAY_SCHEDULE", "opt_learning_rate_decay_schedule"): getattr(args, "schedule_type", "constant") or "constant", getattr(constants, "TENSOR_PARALLELISM", "tensor_parallelism"): train_tp, getattr(constants, "PIPELINE_PARALLELISM", "pipeline_parallelism"): 1, getattr(constants, "CONTEXT_PARALLELISM", "context_parallelism"): train_sp, - getattr(constants, "EXPERT_PARALLELISM", "expert_parallelism"): 1, + getattr(constants, "EXPERT_PARALLELISM", "expert_parallelism"): getattr(args, "train_mesh_expert", 1), "generation_backend": getattr(args, "rollout_engine", "vllm"), "generation_tensor_parallelism": rollout_tp, "generation_pipeline_parallelism": 1, - "generation_expert_parallelism": 1, + "generation_expert_parallelism": getattr(args, "rollout_mesh_expert", 1), getattr( constants, "GENERATION_TRAINING_ROLLOUT_TEMPERATURE", @@ -810,3 +992,4 @@ def init_print( for key, value in logging_configs.items(): if value is not None: mllogger.event(key=key, value=value) + _flush_to_gcs_if_needed()