Add local training logger that writes to {save_dir}/{save_weight}.log

src/trainers/lm/agent.py

@@ -21,7 +21,7 @@ from torch.optim.lr_scheduler import CosineAnnealingLR
 from transformers import AutoTokenizer
 from models import LMConfig, LMForCausalLM
 from dataset import AgentRLDataset
-from utils.training import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
+from utils.training import init_logger, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
 from utils.training import apply_config  # noqa: F401
 from trainers.lm.rollout_engine import create_rollout_engine, compute_per_token_logps
 
@@ -416,6 +416,7 @@ if __name__ == "__main__":
     setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
 
     os.makedirs(args.save_dir, exist_ok=True)
+    init_logger(args.save_dir, getattr(args, "save_weight", "train"))
     lm_config = LMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
                                max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None

src/trainers/lm/distillation.py

@@ -182,6 +182,7 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    init_logger(args.save_dir, getattr(args, "save_weight", "train"))
     lm_config_student = LMConfig(hidden_size=args.student_hidden_size, num_hidden_layers=args.student_num_layers, use_moe=bool(args.student_use_moe))
     lm_config_teacher = LMConfig(hidden_size=args.teacher_hidden_size, num_hidden_layers=args.teacher_num_layers, use_moe=bool(args.teacher_use_moe))
     ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None

src/trainers/lm/dpo.py

@@ -161,6 +161,7 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    init_logger(args.save_dir, getattr(args, "save_weight", "train"))
     lm_config = LMConfig(**vars(args))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     

src/trainers/lm/full_sft.py

@@ -113,6 +113,7 @@ def main(default_config=None):
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    init_logger(args.save_dir, getattr(args, "save_weight", "train"))
     lm_config = LMConfig(**vars(args))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     

src/trainers/lm/grpo.py

@@ -19,7 +19,7 @@ from torch.optim.lr_scheduler import CosineAnnealingLR
 from transformers import AutoModel
 from models import LMConfig, LMForCausalLM
 from dataset import RLAIFDataset
-from utils.training import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
+from utils.training import init_logger, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
 from utils.training import apply_config  # noqa: F401
 from trainers.lm.rollout_engine import create_rollout_engine
 
@@ -249,6 +249,7 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    init_logger(args.save_dir, getattr(args, "save_weight", "train"))
     lm_config = LMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
                                max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None

src/trainers/lm/lora.py

@@ -107,6 +107,7 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    init_logger(args.save_dir, getattr(args, "save_weight", "train"))
     lm_config = LMConfig(**vars(args))
     ckp_data = lm_checkpoint(lm_config, weight=args.lora_name, save_dir='../checkpoints') if args.from_resume==1 else None
     

src/trainers/lm/ppo.py

@@ -18,7 +18,7 @@ from torch.nn.utils import clip_grad_norm_
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from models import LMConfig, LMForCausalLM
 from dataset import RLAIFDataset
-from utils.training import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
+from utils.training import init_logger, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
 from utils.training import apply_config  # noqa: F401
 from trainers.lm.rollout_engine import create_rollout_engine
 
@@ -344,6 +344,7 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    init_logger(args.save_dir, getattr(args, "save_weight", "train"))
     lm_config = LMConfig(**vars(args))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     

src/trainers/lm/pretrain.py

@@ -112,6 +112,7 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    init_logger(args.save_dir, getattr(args, "save_weight", "train"))
     lm_config = LMConfig(**vars(args))
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     

src/trainers/vam/full_sft.py

@@ -170,6 +170,7 @@ def main(default_config=None):
 
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    init_logger(args.save_dir, getattr(args, "save_weight", "train"))
     omni_config = VAMConfig(**vars(args))
     ckp_data = omni_checkpoint(omni_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
 

src/trainers/vlm/full_sft.py

@@ -117,6 +117,7 @@ def main(default_config=None):
 
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    init_logger(args.save_dir, getattr(args, "save_weight", "train"))
     vlm_config = VLMConfig(**vars(args))
     ckp_data = vlm_checkpoint(vlm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
 

src/trainers/vlm/pretrain.py

@@ -117,6 +117,7 @@ if __name__ == "__main__":
 
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
+    init_logger(args.save_dir, getattr(args, "save_weight", "train"))
     vlm_config = VLMConfig(**vars(args))
     ckp_data = vlm_checkpoint(vlm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
 

src/utils/training.py

@@ -57,9 +57,26 @@ def is_main_process():
     return not dist.is_initialized() or dist.get_rank() == 0
 
 
+_log_file = None
+
+
+def init_logger(save_dir='../checkpoint', name='train'):
+    global _log_file
+    if not is_main_process():
+        return
+    os.makedirs(save_dir, exist_ok=True)
+    log_path = os.path.join(save_dir, f'{name}.log')
+    _log_file = open(log_path, 'a', encoding='utf-8')
+    Logger(f'日志写入: {os.path.abspath(log_path)}')
+
+
 def Logger(content):
-    if is_main_process():
-        print(content)
+    if not is_main_process():
+        return
+    print(content)
+    if _log_file is not None:
+        _log_file.write(str(content) + '\n')
+        _log_file.flush()
 
 
 def get_lr(current_step, total_steps, lr):