Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion docs/tutorial/quickstart.md
Original file line number Diff line number Diff line change
Expand Up @@ -97,12 +97,15 @@ python3 training/main_sync_ppo.py --help

## Monitoring the Training Process

We recommend using Weights & Biases (wandb) for monitoring. Run `wandb login` or set the `WANDB_API_KEY` environment variable. Set `wandb.mode=online` in your configuration to upload training statistics.
+ We recommend using [Weights & Biases (wandb)](https://github.com/wandb/wandb) or [SwanLab](https://github.com/SwanHubX/SwanLab) for monitoring—run `wandb login` or `swanlab login`, or set the corresponding environment variable API key (`WANDB_API_KEY` or `SWANLAB_API_KEY`). Set `wandb.mode="online"` or `swanlab.mode="cloud"` in your configuration to upload training statistics. If you cannot connect to the server, you can also use `wandb.mode="offline"` or `swanlab.mode="local"` to save data locally without uploading.


You can also use TensorBoard by setting the `tensorboard.path` parameter.

The main log will be saved to `${fileroot}/logs/${USER}/${experiment_name}/${trial_name}/main.log` and contains the statistics uploaded to wandb.

If SwanLab is enabled, logs will be saved to the directory specified by `swanlab.logdir`.

### Key Training Statistics

- **`Epoch 1/5`**: Indicates the total epochs required and the current epoch being trained.
Expand Down
1 change: 1 addition & 0 deletions evaluation/requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -15,3 +15,4 @@ prettytable
timeout-decorator
timeout_decorator
wandb
swanlab[dashboard]
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@ dependencies = [
"colorlog",
"psutil",
"pynvml",
"swanlab[dashboard]",

# Performance and compression
"ninja",
Expand Down
16 changes: 15 additions & 1 deletion realhf/api/cli_args.py
Original file line number Diff line number Diff line change
Expand Up @@ -848,6 +848,16 @@ class WandBConfig:
config: Optional[Dict] = None


@dataclass
class SwanlabConfig:
project: Optional[str] = None
name: Optional[str] = None
config: Optional[Dict] = None
logdir: Optional[str] = None
mode: Optional[str] = "local"
api_key: Optional[str] = os.getenv("SWANLAB_API_KEY", None)


@dataclass
class TensorBoardConfig:
path: Optional[str] = None
Expand Down Expand Up @@ -951,6 +961,10 @@ class BaseExperimentConfig:
default_factory=WandBConfig,
metadata={"help": "Weights & Biases configuration."},
)
swanlab: SwanlabConfig = field(
default_factory=SwanlabConfig,
metadata={"help": "SwanLab configuration."},
)
tensorboard: TensorBoardConfig = field(
default_factory=TensorBoardConfig,
metadata={"help": "TensorBoard configuration. Only 'path' field required."},
Expand Down Expand Up @@ -1026,7 +1040,7 @@ class BaseExperimentConfig:
default=False,
metadata={
"help": "Enable automatic evaluation during training. "
"Results logged to disk and WandB (if active)."
"Results logged to disk and WandB or Swanlab(if active)."
},
)
auto_eval_config: AutomaticEvaluator = field(
Expand Down
2 changes: 2 additions & 0 deletions realhf/api/core/system_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from realhf.api.cli_args import (
AutomaticEvaluator,
ExperimentSaveEvalControl,
SwanlabConfig,
TensorBoardConfig,
WandBConfig,
)
Expand Down Expand Up @@ -254,6 +255,7 @@ class ExperimentScheduling:
class ExperimentConfig:
exp_ctrl: ExperimentSaveEvalControl
wandb: WandBConfig
swanlab: SwanlabConfig
tensorboard: TensorBoardConfig
# dataflow
model_rpcs: List[dfg.MFCDef]
Expand Down
2 changes: 1 addition & 1 deletion realhf/apps/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ def main_start(args, job_group_id: str = "", recover_count: int = 0):
raise RuntimeError("Experiment initial setup failed.") from e

evaluator = (
AutomaticEvaluator(exp_cfg.evaluator, exp_cfg.wandb)
AutomaticEvaluator(exp_cfg.evaluator, exp_cfg.wandb, exp_cfg.swanlab)
if exp_cfg.auto_eval
else None
)
Expand Down
22 changes: 16 additions & 6 deletions realhf/base/logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,19 +141,29 @@ def getLogger(
return logging.getLogger(name)


_LATEST_WANDB_STEP = 0
_LATEST_LOG_STEP = 0


def log_wandb_tensorboard(data, step=None, summary_writer=None):
import wandb
def log_swanlab_wandb_tensorboard(data, step=None, summary_writer=None):
# Logs data to SwanLab、 wandb、 TensorBoard.

global _LATEST_WANDB_STEP
global _LATEST_LOG_STEP
if step is None:
step = _LATEST_WANDB_STEP
step = _LATEST_LOG_STEP
else:
_LATEST_WANDB_STEP = max(_LATEST_WANDB_STEP, step)
_LATEST_LOG_STEP = max(_LATEST_LOG_STEP, step)

# swanlab
import swanlab

swanlab.log(data, step=step)

# wandb
import wandb

wandb.log(data, step=step)

# tensorboard
if summary_writer is not None:
for key, val in data.items():
summary_writer.add_scalar(f"{key}", val, step)
Expand Down
1 change: 1 addition & 0 deletions realhf/experiments/async_exp/async_rl_exp.py
Original file line number Diff line number Diff line change
Expand Up @@ -331,6 +331,7 @@ def initial_setup(self) -> ExperimentConfig:
return ExperimentConfig(
exp_ctrl=self.exp_ctrl,
wandb=self.wandb,
swanlab=self.swanlab,
tensorboard=self.tensorboard,
# NOTE: master and model worker only see RPCs without generation
model_rpcs=[
Expand Down
1 change: 1 addition & 0 deletions realhf/experiments/common/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -564,6 +564,7 @@ def initial_setup(self) -> ExperimentConfig:
return ExperimentConfig(
exp_ctrl=self.exp_ctrl,
wandb=self.wandb,
swanlab=self.swanlab,
tensorboard=self.tensorboard,
model_rpcs=[rpc_alloc.rpc for rpc_alloc in rpc_allocs],
model_worker=model_worker,
Expand Down
1 change: 1 addition & 0 deletions realhf/experiments/common/ppo_math_exp.py
Original file line number Diff line number Diff line change
Expand Up @@ -370,6 +370,7 @@ def _find_rpc(name):
return ExperimentConfig(
exp_ctrl=self.exp_ctrl,
wandb=self.wandb,
swanlab=self.swanlab,
tensorboard=self.tensorboard,
model_rpcs=[rpc_alloc.rpc for rpc_alloc in rpc_allocs],
model_worker=model_worker,
Expand Down
52 changes: 47 additions & 5 deletions realhf/scheduler/evaluator.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import time
from typing import Dict, Optional

import swanlab
import wandb

import realhf.api.core.system_api as config_pkg
Expand Down Expand Up @@ -125,13 +126,15 @@ def log(self, config: config_pkg.AutomaticEvaluator) -> bool:
self.status = EvaluationStepStatus.FAILED
return False

wandb_data = {}
log_data = {}
for data_name, d in data.items():
for k, v in d.items():
wandb_data[f"{data_name}_{k}"] = v
wandb.log(wandb_data, step=self.global_step)
log_data[f"{data_name}_{k}"] = v
wandb.log(log_data, step=self.global_step)
swanlab.log(log_data, step=self.global_step)
self.status = EvaluationStepStatus.LOGGED
logger.info(f"Logging eval result {wandb_data} to step {self.global_step}")
logger.info(f"Logging eval result {log_data} to step {self.global_step}")

return True

def check(self):
Expand All @@ -154,13 +157,15 @@ def __init__(
self,
config: config_pkg.AutomaticEvaluator,
wandb_config: config_pkg.WandBConfig,
swanlab_config: config_pkg.SwanlabConfig,
):
self.__eval_steps: Dict[int, EvaluationStep] = {}
self.__max_concurrent_jobs = config.max_concurrent_jobs
self.__wandb_config = wandb_config
self.__swanlab_config = swanlab_config
self.__config = config
self.__wandb_initialized = False

self.__swanlab_initialized = False
# Check evaluated checkpoints by logs in recover
# NOTE: All previous evaluation steps with output will be marked
# as logged, even if it is not really logged in wandb.
Expand Down Expand Up @@ -228,6 +233,40 @@ def __lazy_wandb_init(self):
settings=wandb.Settings(start_method="fork"),
)

def __lazy_swanlab_init(self):
if self.__swanlab_config.api_key:
swanlab.login(self.__swanlab_config.api_key)
if self.swanlab_config.config is None:
import yaml

with open(
os.path.join(
constants.LOG_ROOT,
constants.experiment_name(),
constants.trial_name(),
"config.yaml",
),
"r",
) as f:
__config = yaml.safe_load(f)
else:
__config = self.swanlab_config.config
__config["FRAMEWORK"] = "AReaL"
swanlab.init(
project=self.__swanlab_config.project or constants.experiment_name(),
experiment_name=self.__swanlab_config.name
or f"{constants.trial_name()}_eval",
config=__config,
logdir=self.__swanlab_config.logdir
or os.path.join(
constants.LOG_ROOT,
constants.experiment_name(),
constants.trial_name(),
"swanlab",
),
mode=self.__swanlab_config.mode,
)

def step(self):
# Check whether a new evaluation step should be created
ckpt_parent = os.path.join(
Expand Down Expand Up @@ -292,6 +331,9 @@ def step(self):
if not self.__wandb_initialized:
self.__lazy_wandb_init()
self.__wandb_initialized = True
if not self.__swanlab_initialized:
self.__lazy_swanlab_init()
self.__swanlab_initialized = True
self.__eval_steps[log_step].log(self.__config)

@property
Expand Down
38 changes: 37 additions & 1 deletion realhf/system/master_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import colorama
import networkx as nx
import numpy as np
import swanlab
import wandb
from tensorboardX import SummaryWriter

Expand Down Expand Up @@ -312,6 +313,40 @@ def __lazy_init(self):
resume="allow",
settings=wandb.Settings(start_method="fork"),
)

# swanlab init, connect to remote or local swanlab host
if self.swanlab_config.mode != "disabled" and self.swanlab_config.api_key:
swanlab.login(self.swanlab_config.api_key)
if self.swanlab_config.config is None:
import yaml

with open(
os.path.join(
constants.LOG_ROOT,
constants.experiment_name(),
constants.trial_name(),
"config.yaml",
),
"r",
) as f:
__config = yaml.safe_load(f)
else:
__config = self.swanlab_config.config
__config["FRAMEWORK"] = "AReaL"
swanlab.init(
project=self.swanlab_config.project or constants.experiment_name(),
experiment_name=self.swanlab_config.name
or f"{constants.trial_name()}_train",
config=__config,
logdir=self.swanlab_config.logdir
or os.path.join(
constants.LOG_ROOT,
constants.experiment_name(),
constants.trial_name(),
"swanlab",
),
mode=self.swanlab_config.mode,
)
# tensorboard logging
self.__summary_writer = None
if self.tensorboard_config.path is not None:
Expand Down Expand Up @@ -487,7 +522,7 @@ def _log_training_stats(self, e2e_time: float, time_since_configure: float):
s += f"(global step {global_step}) finishes. "
s += f"#End to end# execution time: *{e2e_time:.3f}*s. "
s += f"Total time consumption: {time_since_configure:.3f}s. "
logging.log_wandb_tensorboard({"timeperf/e2e": e2e_time})
logging.log_swanlab_wandb_tensorboard({"timeperf/e2e": e2e_time})
if len(self.e2e_time_history) > 2:
remaining_steps = self._steps_per_epoch - epoch_step
remaining_epochs = self.__total_train_epochs - epoch
Expand Down Expand Up @@ -540,6 +575,7 @@ def experiment_complete_exit(self):
)

wandb.finish()
swanlab.finish()
if self.__summary_writer is not None:
self.__summary_writer.close()
gc.collect()
Expand Down
8 changes: 4 additions & 4 deletions realhf/system/model_function_call.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from collections import defaultdict
from typing import Dict, Hashable, List, Set, Tuple

import swanlab
import wandb
from tensorboardX import SummaryWriter

Expand Down Expand Up @@ -442,7 +443,7 @@ async def run_step(self, buf_indices, sample, buffer_id: int):
logger.info(
f"RPC name {rpc.name} returns\n{data_api.tabulate_stats(res)}"
)
logging.log_wandb_tensorboard(
logging.log_swanlab_wandb_tensorboard(
res,
step=ctrl.step_info.global_step,
summary_writer=self.summary_writer,
Expand All @@ -453,7 +454,7 @@ async def run_step(self, buf_indices, sample, buffer_id: int):
f"RPC name {rpc.name} returns ({j + 1}/{len(res)})\n{data_api.tabulate_stats(r)}"
)
offset = len(res) * ctrl.step_info.global_step
logging.log_wandb_tensorboard(
logging.log_swanlab_wandb_tensorboard(
r,
step=offset + j,
summary_writer=self.summary_writer,
Expand All @@ -465,11 +466,10 @@ async def run_step(self, buf_indices, sample, buffer_id: int):
for time_record in time_records:
stats_tracker.scalar(**time_record)
time_stats = stats_tracker.export()
logging.log_wandb_tensorboard(
logging.log_swanlab_wandb_tensorboard(
time_stats,
summary_writer=self.summary_writer,
)

logger.info(
f"Model rpc {rpc.name} finished. "
f"Request-reply time {time.perf_counter() - tik:.4f}s. "
Expand Down
2 changes: 2 additions & 0 deletions realhf/system/worker_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -580,7 +580,9 @@ def configure(
)
expr_config.lazy_init()
self.wandb_config = expr_config.wandb
self.swanlab_config = expr_config.swanlab
os.environ["WANDB_MODE"] = self.wandb_config.mode
os.environ["SWANLAB_MODE"] = self.swanlab_config.mode
self.tensorboard_config = expr_config.tensorboard
config = expr_config.resolve_worker_config(
self.__worker_type, self.__worker_index
Expand Down
3 changes: 2 additions & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -68,4 +68,5 @@ python_dateutil
word2number
Pebble
timeout-decorator
prettytable
prettytable
swanlab[dashboard]
7 changes: 7 additions & 0 deletions training/configs/async-ppo.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,13 @@ wandb:
notes: null
tags: null
config: null
swanlab:
mode: disabled
api_key: null
project: null
name: null
config: null
logdir: null
tensorboard:
path: null
recover_mode: auto
Expand Down
Loading