Integrate SwanLab for training log tracking

README.md

@@ -110,7 +110,9 @@
 - 收集、蒸馏、整理并清洗去重所有阶段的高质量数据集，且全部开源。
 - 从0实现预训练、指令微调、LoRA、DPO强化学习，白盒模型蒸馏。关键算法几乎不依赖第三方封装的框架，且全部开源。
 - 同时兼容`transformers`、`trl`、`peft`等第三方主流框架。
-- 训练支持单机单卡、单机多卡(DDP、DeepSpeed)训练，支持wandb可视化训练流程。支持动态启停训练。
+
+* 训练支持单机单卡、单机多卡(DDP、DeepSpeed)训练，支持wandb、swanlab可视化训练流程。支持动态启停训练。
+
 - 在第三方测评榜（C-Eval、C-MMLU、OpenBookQA等）进行模型测试。
 - 实现Openai-Api协议的极简服务端，便于集成到第三方ChatUI使用（FastGPT、Open-WebUI等）。
 - 基于streamlit实现最简聊天WebUI前端。
@@ -339,17 +341,20 @@ torchrun --nproc_per_node N train_xxx.py
 deepspeed --master_port 29500 --num_gpus=N train_xxx.py
 ```
 
-可根据需要开启wandb记录训练过程
+可根据需要开启wandb或者swanlab记录训练过程
 
 ```bash
-# 需要登录: wandb login
+# 以wandb为例，需要登录: wandb login
-torchrun --nproc_per_node N train_xxx.py --use_wandb
+torchrun --nproc_per_node N train_xxx.py --report_to wandb
 # and
-python train_xxx.py --use_wandb
+python train_xxx.py --report_to wandb
 ```
 
-通过添加`--use_wandb`参数，可以记录训练过程，训练完成后，可以在wandb网站上查看训练过程。通过修改`wandb_project`
+通过添加`--report_to <wandb或者swanlab>`参数，可以使用在线跟踪工具记录训练过程，训练完成后，可以在[wandb网站](https://wandb.ai)或者
-和`wandb_run_name`参数，可以指定项目名称和运行名称。
+[swanlab网站](https://swanlab.cn)上查看训练过程。通过添加`--project_name <项目名称>`和`--run_name <实验名称>`参数，可以指定项目名称和运行名称。
+
+如果训练服务器无法访问外网，你可以通过添加参数`--report_to swanlab`并根据引导选择3开启离线日志看板模式。在终端使用命令
+`swanlab watch -h 0.0.0.0 -p 8080` 启动SwanLab离线仪表板。
 
 </details>
 

README_en.md

@@ -369,14 +369,19 @@ Enable wandb to record the training process if needed:
 
 ```bash
 # Need to log in: wandb login
-torchrun --nproc_per_node N train_xxx.py --use_wandb
+torchrun --nproc_per_node N train_xxx.py --report_to wandb
 # and
-python train_xxx.py --use_wandb
+python train_xxx.py --report_to wandb
 ```
 
-By adding the `--use_wandb` parameter, the training process will be recorded, and after training, you can view the
+By adding the `--report_to wandb` parameter, the training process will be recorded, and after training, you can view the
-process on the wandb website. Modify the `wandb_project` and `wandb_run_name` parameters to specify project and run
+process on the [wandb](https://wandb.ai) website. Modify the `--project_name <PROJECT NAME>` and `--run_name <RUN NAME>` parameters
-names.
+to specify project and run names.
+
+If you prefer using SwanLab or if the server cannot access the internet, you are welcome to use SwanLab and its
+[offline mode](https://docs.swanlab.cn/en/guide_cloud/self_host/offline-board.html) by adding the parameter `--report_to swanlab` and
+following the instruct in terminal. Then open the terminal and use `swanlab watch -h 0.0.0.0 -p 8080` command to start a SwanLab
+offline dashboard.
 
 </details>
 

train_distill_reason.py

@@ -27,11 +27,27 @@ def Logger(content):
         print(content)
 
 
+def init_tracker(args):
+    if args.report_to == "wandb":
+        import wandb
+
+        wandb.init(project=args.project_name, run_name=args.run_name)
+        tracker = wandb
+    if args.report_to == "swanlab":
+        import swanlab
+    
+        swanlab.init(project=args.project_name, run_name=args.run_name, config=args)
+        tracker = swanlab
+    else:
+        tracker = None
+    return tracker
+
+
 def get_lr(current_step, total_steps, lr):
     return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
 
 
-def train_epoch(epoch, wandb):
+def train_epoch(epoch, tracker):
     # 思考标签占位符
     start_of_think_ids = tokenizer('<think>').input_ids
     end_of_think_ids = tokenizer('</think>').input_ids
@@ -40,9 +56,9 @@ def train_epoch(epoch, wandb):
     loss_fct = nn.CrossEntropyLoss(reduction='none')
     start_time = time.time()
     for step, (X, Y, loss_mask) in enumerate(train_loader):
-        X = X.to(args.device)
+        X = X.to(DEVICE)
-        Y = Y.to(args.device)
+        Y = Y.to(DEVICE)
-        loss_mask = loss_mask.to(args.device)
+        loss_mask = loss_mask.to(DEVICE)
         lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
@@ -56,7 +72,7 @@ def train_epoch(epoch, wandb):
             sp_ids = torch.isin(Y.view(-1),
                                 torch.tensor(start_of_think_ids + end_of_think_ids
                                              + start_of_answer_ids + end_of_answer_ids
-                                             ).to(args.device))
+                                             ).to(DEVICE))
             # 在 sp_ids 对应的位置增加额外的惩罚
             loss_mask = loss_mask.view(-1)
             loss_mask_sum = loss_mask.sum()
@@ -89,8 +105,8 @@ def train_epoch(epoch, wandb):
                     optimizer.param_groups[-1]['lr'],
                     spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
 
-            if (wandb is not None) and (not ddp or dist.get_rank() == 0):
+            if (tracker is not None) and (not ddp or dist.get_rank() == 0):
-                wandb.log({"loss": loss,
+                tracker.log({"loss": loss,
                            "lr": optimizer.param_groups[-1]['lr'],
                            "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
 
@@ -113,10 +129,10 @@ def init_model(lm_config):
     model = MiniMindLM(lm_config)
     moe_path = '_moe' if lm_config.use_moe else ''
     ckp = f'./out/rlhf_{lm_config.dim}{moe_path}.pth'
-    state_dict = torch.load(ckp, map_location=args.device)
+    state_dict = torch.load(ckp, map_location=DEVICE)
     model.load_state_dict(state_dict, strict=False)
     Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
-    model = model.to(args.device)
+    model = model.to(DEVICE)
     return model, tokenizer
 
 
@@ -140,8 +156,8 @@ if __name__ == "__main__":
     parser.add_argument("--learning_rate", type=float, default=1e-6)
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
     parser.add_argument("--dtype", type=str, default="bfloat16")
-    parser.add_argument("--use_wandb", action="store_true")
+    parser.add_argument("--report_to", type=str, default="")
-    parser.add_argument("--wandb_project", type=str, default="MiniMind-Full-SFT")
+    parser.add_argument("--project_name", type=str, default="MiniMind-Full-SFT")
     parser.add_argument("--num_workers", type=int, default=1)
     parser.add_argument("--ddp", action="store_true")
     parser.add_argument("--accumulation_steps", type=int, default=1)
@@ -166,21 +182,19 @@ if __name__ == "__main__":
     torch.manual_seed(1337)
     device_type = "cuda" if "cuda" in args.device else "cpu"
 
-    args.wandb_run_name = f"MiniMind-Distill-Reasoning-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
+    args.run_name = f"MiniMind-Distill-Reasoning-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
 
     ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
     ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
     ddp_local_rank, DEVICE = 0, "cuda:0"
     if ddp:
         init_distributed_mode()
-        args.device = torch.device(DEVICE)
+        args.device = DEVICE
 
-    if args.use_wandb and (not ddp or ddp_local_rank == 0):
+    if not ddp or ddp_local_rank == 0:
-        import wandb
+        tracker = init_tracker(args)
-
-        wandb.init(project=args.wandb_project, name=args.wandb_run_name)
     else:
-        wandb = None
+        tracker = None
 
     model, tokenizer = init_model(lm_config)
 
@@ -205,4 +219,4 @@ if __name__ == "__main__":
 
     iter_per_epoch = len(train_loader)
     for epoch in range(args.epochs):
-        train_epoch(epoch, wandb)
+        train_epoch(epoch, tracker)

train_distillation.py

@@ -26,6 +26,22 @@ def Logger(content):
         print(content)
 
 
+def init_tracker(args):
+    if args.report_to == "wandb":
+        import wandb
+
+        wandb.init(project=args.project_name, run_name=args.run_name)
+        tracker = wandb
+    if args.report_to == "swanlab":
+        import swanlab
+    
+        swanlab.init(project=args.project_name, run_name=args.run_name, config=args)
+        tracker = swanlab
+    else:
+        tracker = None
+    return tracker
+
+
 def get_lr(current_step, total_steps, lr):
     return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
 
@@ -44,7 +60,7 @@ def distillation_loss_fn(student_logits, teacher_logits, temperature=1.0, reduct
     return (temperature ** 2) * kl
 
 
-def train_epoch(epoch, wandb, alpha=0.0, temperature=1.0):
+def train_epoch(epoch, tracker, alpha=0.0, temperature=1.0):
     start_time = time.time()
 
     if teacher_model is not None:
@@ -52,9 +68,9 @@ def train_epoch(epoch, wandb, alpha=0.0, temperature=1.0):
         teacher_model.requires_grad_(False)
 
     for step, (X, Y, loss_mask) in enumerate(train_loader):
-        X = X.to(args.device)
+        X = X.to(DEVICE)
-        Y = Y.to(args.device)
+        Y = Y.to(DEVICE)
-        loss_mask = loss_mask.to(args.device)
+        loss_mask = loss_mask.to(DEVICE)
         lr = get_lr(epoch * iter_per_epoch + step,
                     args.epochs * iter_per_epoch,
                     args.learning_rate)
@@ -95,7 +111,7 @@ def train_epoch(epoch, wandb, alpha=0.0, temperature=1.0):
                 temperature=temperature
             )
         else:
-            distill_loss = torch.tensor(0.0, device=args.device)
+            distill_loss = torch.tensor(0.0, device=DEVICE)
 
         # 3) 总损失 = alpha * CE + (1-alpha) * Distill
         loss = alpha * ce_loss + (1 - alpha) * distill_loss
@@ -123,8 +139,8 @@ def train_epoch(epoch, wandb, alpha=0.0, temperature=1.0):
                 )
             )
 
-            if (wandb is not None) and (not ddp or dist.get_rank() == 0):
+            if (tracker is not None) and (not ddp or dist.get_rank() == 0):
-                wandb.log({
+                tracker.log({
                     "loss": loss.item(),
                     "ce_loss": ce_loss.item(),
                     "distill_loss": distill_loss.item() if teacher_model is not None else 0.0,
@@ -149,10 +165,10 @@ def init_student_model(lm_config):
     model = MiniMindLM(lm_config)
     moe_path = '_moe' if lm_config.use_moe else ''
     ckp = f'./out/full_sft_{lm_config.dim}{moe_path}.pth'
-    state_dict = torch.load(ckp, map_location=args.device)
+    state_dict = torch.load(ckp, map_location=DEVICE)
     model.load_state_dict(state_dict, strict=False)
     Logger(f'学生模型(LLM)总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
-    model = model.to(args.device)
+    model = model.to(DEVICE)
 
     return model, tokenizer
 
@@ -161,10 +177,10 @@ def init_teacher_model(lm_config):
     model = MiniMindLM(lm_config)
     moe_path = '_moe' if lm_config.use_moe else ''
     ckp = f'./out/full_sft_{lm_config.dim}{moe_path}.pth'
-    state_dict = torch.load(ckp, map_location=args.device)
+    state_dict = torch.load(ckp, map_location=DEVICE)
     model.load_state_dict(state_dict, strict=False)
     Logger(f'教师模型(LLM)总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
-    model = model.to(args.device)
+    model = model.to(DEVICE)
     return model
 
 
@@ -188,8 +204,8 @@ if __name__ == "__main__":
     parser.add_argument("--learning_rate", type=float, default=5e-6)
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
     parser.add_argument("--dtype", type=str, default="bfloat16")
-    parser.add_argument("--use_wandb", action="store_true")
+    parser.add_argument("--report_to", type=str, default="")
-    parser.add_argument("--wandb_project", type=str, default="MiniMind-Full-SFT")
+    parser.add_argument("--project_name", type=str, default="MiniMind-Full-SFT")
     parser.add_argument("--num_workers", type=int, default=1)
     parser.add_argument("--ddp", action="store_true")
     parser.add_argument("--accumulation_steps", type=int, default=1)
@@ -212,21 +228,19 @@ if __name__ == "__main__":
     torch.manual_seed(1337)
     device_type = "cuda" if "cuda" in args.device else "cpu"
 
-    args.wandb_run_name = f"MiniMind-Dist-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
+    args.un_name = f"MiniMind-Dist-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
 
     ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
     ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
     ddp_local_rank, DEVICE = 0, "cuda:0"
     if ddp:
         init_distributed_mode()
-        args.device = torch.device(DEVICE)
+        args.device = DEVICE
 
-    if args.use_wandb and (not ddp or ddp_local_rank == 0):
+    if not ddp or ddp_local_rank == 0:
-        import wandb
+        tracker = init_tracker(args)
-
-        wandb.init(project=args.wandb_project, name=args.wandb_run_name)
     else:
-        wandb = None
+        tracker = None
 
     # 初始化学生模型和教师模型
     model, tokenizer = init_student_model(lm_config_student)
@@ -253,4 +267,4 @@ if __name__ == "__main__":
 
     iter_per_epoch = len(train_loader)
     for epoch in range(args.epochs):
-        train_epoch(epoch, wandb)
+        train_epoch(epoch, tracker)

train_dpo.py

@@ -27,6 +27,22 @@ def Logger(content):
         print(content)
 
 
+def init_tracker(args):
+    if args.report_to == "wandb":
+        import wandb
+
+        wandb.init(project=args.project_name, run_name=args.run_name)
+        tracker = wandb
+    if args.report_to == "swanlab":
+        import swanlab
+    
+        swanlab.init(project=args.project_name, run_name=args.run_name, config=args)
+        tracker = swanlab
+    else:
+        tracker = None
+    return tracker
+
+
 def get_lr(current_step, total_steps, lr):
     return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
 
@@ -60,15 +76,15 @@ def dpo_loss(ref_probs, probs, beta):
     return loss.mean()
 
 
-def train_epoch(epoch, wandb):
+def train_epoch(epoch, tracker):
     start_time = time.time()
     for step, batch in enumerate(train_loader):
-        x_chosen = batch['x_chosen'].to(args.device)
+        x_chosen = batch['x_chosen'].to(DEVICE)
-        x_rejected = batch['x_rejected'].to(args.device)
+        x_rejected = batch['x_rejected'].to(DEVICE)
-        y_chosen = batch['y_chosen'].to(args.device)
+        y_chosen = batch['y_chosen'].to(DEVICE)
-        y_rejected = batch['y_rejected'].to(args.device)
+        y_rejected = batch['y_rejected'].to(DEVICE)
-        mask_chosen = batch['mask_chosen'].to(args.device)
+        mask_chosen = batch['mask_chosen'].to(DEVICE)
-        mask_rejected = batch['mask_rejected'].to(args.device)
+        mask_rejected = batch['mask_rejected'].to(DEVICE)
         x = torch.cat([x_chosen, x_rejected], dim=0)
         y = torch.cat([y_chosen, y_rejected], dim=0)
         mask = torch.cat([mask_chosen, mask_rejected], dim=0)
@@ -111,8 +127,8 @@ def train_epoch(epoch, wandb):
                     optimizer.param_groups[-1]['lr'],
                     spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
 
-            if (wandb is not None) and (not ddp or dist.get_rank() == 0):
+            if (tracker is not None) and (not ddp or dist.get_rank() == 0):
-                wandb.log({"loss": loss,
+                tracker.log({"loss": loss,
                            "lr": optimizer.param_groups[-1]['lr'],
                            "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
 
@@ -135,7 +151,7 @@ def init_model(lm_config):
     model = MiniMindLM(lm_config)
     moe_path = '_moe' if lm_config.use_moe else ''
     ckp = f'./out/full_sft_{lm_config.dim}{moe_path}.pth'
-    state_dict = torch.load(ckp, map_location=args.device)
+    state_dict = torch.load(ckp, map_location=DEVICE)
     model.load_state_dict(state_dict, strict=False)
     # 初始化参考模型
     ref_model = MiniMindLM(lm_config)
@@ -144,8 +160,8 @@ def init_model(lm_config):
     ref_model.requires_grad_(False)
 
     Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
-    model = model.to(args.device)
+    model = model.to(DEVICE)
-    ref_model = ref_model.to(args.device)
+    ref_model = ref_model.to(DEVICE)
 
     return model, ref_model, tokenizer
 
@@ -171,8 +187,8 @@ if __name__ == "__main__":
     parser.add_argument("--learning_rate", type=float, default=1e-8)
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
     parser.add_argument("--dtype", type=str, default="bfloat16")
-    parser.add_argument("--use_wandb", action="store_true")
+    parser.add_argument("--report_to", type=str, default="")
-    parser.add_argument("--wandb_project", type=str, default="MiniMind-RLHF-SFT")
+    parser.add_argument("--project_name", type=str, default="MiniMind-RLHF-SFT")
     parser.add_argument("--num_workers", type=int, default=1)
     parser.add_argument("--ddp", action="store_true")
     parser.add_argument("--accumulation_steps", type=int, default=1)
@@ -197,21 +213,19 @@ if __name__ == "__main__":
     torch.manual_seed(1337)
     device_type = "cuda" if "cuda" in args.device else "cpu"
 
-    args.wandb_run_name = f"MiniMind-Full-DPO-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
+    args.run_name = f"MiniMind-Full-DPO-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
 
     ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
     ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
     ddp_local_rank, DEVICE = 0, "cuda:0"
     if ddp:
         init_distributed_mode()
-        args.device = torch.device(DEVICE)
+        args.device = DEVICE
 
-    if args.use_wandb and (not ddp or ddp_local_rank == 0):
+    if not ddp or ddp_local_rank == 0:
-        import wandb
+        tracker = init_tracker(args)
-
-        wandb.init(project=args.wandb_project, name=args.wandb_run_name)
     else:
-        wandb = None
+        tracker = None
 
     model, ref_model, tokenizer = init_model(lm_config)
 
@@ -236,4 +250,4 @@ if __name__ == "__main__":
 
     iter_per_epoch = len(train_loader)
     for epoch in range(args.epochs):
-        train_epoch(epoch, wandb)
+        train_epoch(epoch, tracker)

train_full_sft.py

@@ -27,17 +27,34 @@ def Logger(content):
         print(content)
 
 
+def init_tracker(args):
+    if args.report_to == "wandb":
+        import wandb
+
+        wandb.init(project=args.project_name, run_name=args.run_name)
+        tracker = wandb
+    if args.report_to == "swanlab":
+        import swanlab
+    
+        swanlab.init(project=args.project_name, run_name=args.run_name, config=args)
+        tracker = swanlab
+    else:
+        tracker = None
+    return tracker
+
+
+
 def get_lr(current_step, total_steps, lr):
     return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
 
 
-def train_epoch(epoch, wandb):
+def train_epoch(epoch, tracker):
     loss_fct = nn.CrossEntropyLoss(reduction='none')
     start_time = time.time()
     for step, (X, Y, loss_mask) in enumerate(train_loader):
-        X = X.to(args.device)
+        X = X.to(DEVICE)
-        Y = Y.to(args.device)
+        Y = Y.to(DEVICE)
-        loss_mask = loss_mask.to(args.device)
+        loss_mask = loss_mask.to(DEVICE)
         lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
@@ -76,8 +93,8 @@ def train_epoch(epoch, wandb):
                     optimizer.param_groups[-1]['lr'],
                     spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
 
-            if (wandb is not None) and (not ddp or dist.get_rank() == 0):
+            if (tracker is not None) and (not ddp or dist.get_rank() == 0):
-                wandb.log({"loss": loss,
+                tracker.log({"loss": loss,
                            "lr": optimizer.param_groups[-1]['lr'],
                            "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
 
@@ -97,7 +114,7 @@ def train_epoch(epoch, wandb):
 
 def init_model(lm_config):
     tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
-    model = MiniMindLM(lm_config)
+    model = MiniMindLM(lm_config).to(DEVICE)
     moe_path = '_moe' if lm_config.use_moe else ''
     ckp = f'./out/pretrain_{lm_config.dim}{moe_path}.pth'
     state_dict = torch.load(ckp, map_location=args.device)
@@ -127,8 +144,8 @@ if __name__ == "__main__":
     parser.add_argument("--learning_rate", type=float, default=5e-5)
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
     parser.add_argument("--dtype", type=str, default="bfloat16")
-    parser.add_argument("--use_wandb", action="store_true")
+    parser.add_argument("--report_to", type=str, default="")
-    parser.add_argument("--wandb_project", type=str, default="MiniMind-Full-SFT")
+    parser.add_argument("--project_name", type=str, default="MiniMind-Full-SFT")
     parser.add_argument("--num_workers", type=int, default=1)
     parser.add_argument("--ddp", action="store_true")
     parser.add_argument("--accumulation_steps", type=int, default=1)
@@ -153,21 +170,19 @@ if __name__ == "__main__":
     torch.manual_seed(1337)
     device_type = "cuda" if "cuda" in args.device else "cpu"
 
-    args.wandb_run_name = f"MiniMind-Full-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
+    args.run_name = f"MiniMind-Full-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
 
     ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
     ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
     ddp_local_rank, DEVICE = 0, "cuda:0"
     if ddp:
         init_distributed_mode()
-        args.device = torch.device(DEVICE)
+        args.device = DEVICE
 
-    if args.use_wandb and (not ddp or ddp_local_rank == 0):
+    if not ddp or ddp_local_rank == 0:
-        import wandb
+        tracker = init_tracker(args)
-
-        wandb.init(project=args.wandb_project, name=args.wandb_run_name)
     else:
-        wandb = None
+        tracker = None
 
     model, tokenizer = init_model(lm_config)
 
@@ -192,4 +207,4 @@ if __name__ == "__main__":
 
     iter_per_epoch = len(train_loader)
     for epoch in range(args.epochs):
-        train_epoch(epoch, wandb)
+        train_epoch(epoch, tracker)

train_lora.py

@@ -23,18 +23,34 @@ def Logger(content):
         print(content)
 
 
+def init_tracker(args):
+    if args.report_to == "wandb":
+        import wandb
+
+        wandb.init(project=args.project_name, run_name=args.run_name)
+        tracker = wandb
+    if args.report_to == "swanlab":
+        import swanlab
+    
+        swanlab.init(project=args.project_name, run_name=args.run_name, config=args)
+        tracker = swanlab
+    else:
+        tracker = None
+    return tracker
+
+
 def get_lr(current_step, total_steps, lr):
     return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
 
 
 # 代码和full_sft「几乎」一致
-def train_epoch(epoch, wandb):
+def train_epoch(epoch, tracker):
     loss_fct = nn.CrossEntropyLoss(reduction='none')
     start_time = time.time()
     for step, (X, Y, loss_mask) in enumerate(train_loader):
-        X = X.to(args.device)
+        X = X.to(DEVICE)
-        Y = Y.to(args.device)
+        Y = Y.to(DEVICE)
-        loss_mask = loss_mask.to(args.device)
+        loss_mask = loss_mask.to(DEVICE)
         lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
@@ -72,8 +88,8 @@ def train_epoch(epoch, wandb):
                     optimizer.param_groups[-1]['lr'],
                     spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
 
-            if (wandb is not None) and (not ddp or dist.get_rank() == 0):
+            if (tracker is not None) and (not ddp or dist.get_rank() == 0):
-                wandb.log({"loss": loss,
+                tracker.log({"loss": loss,
                            "lr": optimizer.param_groups[-1]['lr'],
                            "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
 
@@ -89,9 +105,9 @@ def init_model(lm_config):
     model = MiniMindLM(lm_config)
     moe_path = '_moe' if lm_config.use_moe else ''
     ckp = f'./out/rlhf_{lm_config.dim}{moe_path}.pth'
-    state_dict = torch.load(ckp, map_location=args.device)
+    state_dict = torch.load(ckp, map_location=DEVICE)
     model.load_state_dict(state_dict, strict=False)
-    return model.to(args.device), tokenizer
+    return model.to(DEVICE), tokenizer
 
 
 def init_distributed_mode():
@@ -114,8 +130,8 @@ if __name__ == "__main__":
     parser.add_argument("--learning_rate", type=float, default=5e-5)
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
     parser.add_argument("--dtype", type=str, default="bfloat16")
-    parser.add_argument("--use_wandb", action="store_true")
+    parser.add_argument("--report_to", type=str, default="")
-    parser.add_argument("--wandb_project", type=str, default="MiniMind-LoRA-SFT")
+    parser.add_argument("--project_name", type=str, default="MiniMind-LoRA-SFT")
     parser.add_argument("--num_workers", type=int, default=1)
     parser.add_argument("--ddp", action="store_true")
     parser.add_argument("--accumulation_steps", type=int, default=1)
@@ -145,15 +161,13 @@ if __name__ == "__main__":
     ddp_local_rank, DEVICE = 0, "cuda:0"
     if ddp:
         init_distributed_mode()
-        args.device = torch.device(DEVICE)
+        args.device = DEVICE
 
-    args.wandb_run_name = f"MiniMind-Lora-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
+    args.run_name = f"MiniMind-Lora-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
-    if args.use_wandb and (not ddp or ddp_local_rank == 0):
+    if not ddp or ddp_local_rank == 0:
-        import wandb
+        tracker = init_tracker(args)
-
-        wandb.init(project=args.wandb_project, name=args.wandb_run_name)
     else:
-        wandb = None
+        tracker = None
 
     model, tokenizer = init_model(lm_config)
     apply_lora(model)
@@ -191,4 +205,4 @@ if __name__ == "__main__":
     iter_per_epoch = len(train_loader)
 
     for epoch in range(args.epochs):
-        train_epoch(epoch, wandb)
+        train_epoch(epoch, tracker)

train_pretrain.py

@@ -27,17 +27,33 @@ def Logger(content):
         print(content)
 
 
+def init_tracker(args):
+    if args.report_to == "wandb":
+        import wandb
+
+        wandb.init(project=args.project_name, run_name=args.run_name)
+        tracker = wandb
+    if "swanlab" in args.report_to:
+        import swanlab
+
+        swanlab.init(project=args.project_name, run_name=args.run_name, config=args)
+        tracker = swanlab
+    else:
+        tracker = None
+    return tracker
+
+
 def get_lr(current_step, total_steps, lr):
     return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
 
 
-def train_epoch(epoch, wandb):
+def train_epoch(epoch, tracker):
     loss_fct = nn.CrossEntropyLoss(reduction='none')
     start_time = time.time()
     for step, (X, Y, loss_mask) in enumerate(train_loader):
-        X = X.to(args.device)
+        X = X.to(DEVICE)
-        Y = Y.to(args.device)
+        Y = Y.to(DEVICE)
-        loss_mask = loss_mask.to(args.device)
+        loss_mask = loss_mask.to(DEVICE)
 
         lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
         for param_group in optimizer.param_groups:
@@ -76,8 +92,8 @@ def train_epoch(epoch, wandb):
                     optimizer.param_groups[-1]['lr'],
                     spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
 
-            if (wandb is not None) and (not ddp or dist.get_rank() == 0):
+            if (tracker is not None) and (not ddp or dist.get_rank() == 0):
-                wandb.log({"loss": loss.item() * args.accumulation_steps,
+                tracker.log({"loss": loss,
                            "lr": optimizer.param_groups[-1]['lr'],
                            "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
 
@@ -97,7 +113,7 @@ def train_epoch(epoch, wandb):
 
 def init_model(lm_config):
     tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
-    model = MiniMindLM(lm_config).to(args.device)
+    model = MiniMindLM(lm_config).to(DEVICE)
     Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
     return model, tokenizer
 
@@ -124,8 +140,8 @@ if __name__ == "__main__":
     parser.add_argument("--learning_rate", type=float, default=5e-4)
     parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
     parser.add_argument("--dtype", type=str, default="bfloat16")
-    parser.add_argument("--use_wandb", action="store_true")
+    parser.add_argument("--report_to", type=str, default="")
-    parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain")
+    parser.add_argument("--project_name", type=str, default="MiniMind-Pretrain")
     parser.add_argument("--num_workers", type=int, default=1)
     parser.add_argument("--ddp", action="store_true")
     parser.add_argument("--accumulation_steps", type=int, default=8)
@@ -149,7 +165,7 @@ if __name__ == "__main__":
     torch.manual_seed(1337)
     device_type = "cuda" if "cuda" in args.device else "cpu"
 
-    args.wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
+    args.run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
 
     ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
 
@@ -158,14 +174,12 @@ if __name__ == "__main__":
 
     if ddp:
         init_distributed_mode()
-        args.device = torch.device(DEVICE)
+        args.device = DEVICE
 
-    if args.use_wandb and (not ddp or ddp_local_rank == 0):
+    if not ddp or ddp_local_rank == 0:
-        import wandb
+        tracker = init_tracker(args)
-
-        wandb.init(project=args.wandb_project, name=args.wandb_run_name)
     else:
-        wandb = None
+        tracker = None
 
     model, tokenizer = init_model(lm_config)
     train_ds = PretrainDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
@@ -189,4 +203,4 @@ if __name__ == "__main__":
 
     iter_per_epoch = len(train_loader)
     for epoch in range(args.epochs):
-        train_epoch(epoch, wandb)
+        train_epoch(epoch, tracker)