polish(pu): add stable_adaptor_scale

Author: puyuan

lzero/entry/train_unizero_multitask_balance_segment_ddp.py

@@ -32,6 +32,19 @@
 from collections import defaultdict
 GLOBAL_EVAL_RETURNS: dict[int, float] = defaultdict(lambda: None)
 
+def log_param_statistics(model, logger=logging):
+    n_tensors_total   = sum(1 for _ in model.parameters())
+    n_tensors_train   = sum(p.requires_grad for p in model.parameters())
+
+    n_elems_total     = sum(p.numel() for p in model.parameters())
+    n_elems_train     = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    logger.info(
+        f'Trainable parameters: '
+        f'{n_tensors_train}/{n_tensors_total} tensors  |  '
+        f'{n_elems_train:,}/{n_elems_total:,} elements '
+        f'(~{n_elems_train/1e6:.2f} M / {n_elems_total/1e6:.2f} M)'
+    )
 
 def tasks_per_stage(unsolved: int, remain_lora: int) -> int:
     """
@@ -84,6 +97,10 @@ def step(self, solved_cnt: int, unsolved_cnt: int, train_iter: int):
             logging.info(f'[Curriculum] switch to stage {self.stage} '
                          f'(solved={solved_cnt}, unsolved={unsolved_cnt}, '
                          f'iter={train_iter})')
+                        
+            updated = sum(p.requires_grad for p in self.policy._learn_model.world_model.parameters())
+            logging.info(f'{updated}/{sum(1 for _ in self.policy._learn_model.world_model.parameters())} params will be optimized')
+            log_param_statistics(self.policy._learn_model.world_model)          # 再打印一次，看看数值变化
             self.last_solved      = solved_cnt
             self.last_switch_iter = train_iter
             return True
@@ -595,6 +612,9 @@ def train_unizero_multitask_balance_segment_ddp(
     # 初始化一次（rank0 或各 rank 均可）
     curr_ctrl = CurriculumController(cfg, policy)
 
+    updated = sum(p.requires_grad for p in policy._learn_model.world_model.parameters())
+    logging.info(f'{updated}/{sum(1 for _ in policy._learn_model.world_model.parameters())} params will be optimized')
+    
     while True:
         last_curriculum_stage = cur_curriculum_stage
 
@@ -814,9 +834,14 @@ def train_unizero_multitask_balance_segment_ddp(
             for module_name, module in transformer.named_modules():
                 if isinstance(module, CurriculumLoRALinear) and module.adapters is not None:
                     for adapter_idx, scale_param in enumerate(module.adapter_scales):
+                        # tb_logger.add_scalar(
+                        #     f'UniZero-MT/adapter_scales/{module_name}/adapter_{adapter_idx}',
+                        #     scale_param.item(),
+                        #     global_step=learner.train_iter
+                        # )
                         tb_logger.add_scalar(
                             f'UniZero-MT/adapter_scales/{module_name}/adapter_{adapter_idx}',
-                            scale_param.item(),
+                            scale_param().item(),
                             global_step=learner.train_iter
                         )
 

lzero/model/unizero_world_models/transformer.py

@@ -23,6 +23,21 @@
 from lzero.model.common import SimNorm
 import logging
 
+class LearnableScale(nn.Module):
+    """
+    可学习且有界的标量参数:
+      s = s_max * sigmoid(ŝ)         (0, s_max)
+    """
+    def __init__(self, init=1.0, s_max=1.5):
+        super().__init__()
+        # 反推初始值
+        inv_sig = math.log(init / (s_max - init + 1e-9))
+        self.logit = nn.Parameter(torch.tensor(inv_sig))
+        self.logit.requires_grad = True # TODO
+        self.s_max = s_max
+
+    def forward(self):
+        return self.s_max * torch.sigmoid(self.logit)
 ##############################################
 # CurriculumLoRALinear 实现
 ##############################################
@@ -74,7 +89,9 @@ def __init__(self, in_features: int, out_features: int, bias: bool = True,
 
         # 初始化 LoRA adapter，只有在 r > 0 且 curriculum_stage_num > 1 时才存在
         self.adapters = nn.ModuleList()
-        self.adapter_scales = nn.ParameterList()
+        # self.adapter_scales = nn.ParameterList()
+        self.adapter_scales = nn.ModuleList()
+
         if r > 0 and (curriculum_stage_num - 1) > 0:
             for i in range(curriculum_stage_num - 1):
                 adapter = nn.ParameterDict({
@@ -83,9 +100,15 @@ def __init__(self, in_features: int, out_features: int, bias: bool = True,
                 })
                 self.adapters.append(adapter)
 
-                self.adapter_scales.append(  #  ← 新增
-                    nn.Parameter(torch.tensor(lora_scale_init, dtype=torch.float32))
-                )
+                self.adapter_scales.append(LearnableScale(lora_scale_init, s_max=1.5))
+                
+                # self.adapter_scales.append(  #  ← 新增
+                #     nn.Parameter(torch.tensor(lora_scale_init, dtype=torch.float32))
+                # )
+
+            # --- CurriculumLoRALinear.__init__() ------------
+            # for p in self.adapter_scales:
+            #     p.requires_grad = True   # 统一设 True，避免遗漏
         else:
             self.adapters = None
 
@@ -121,17 +144,21 @@ def set_curriculum_stage(self, stage: int):
                 for idx, adapter in enumerate(self.adapters):
                     adapter['lora_A'].requires_grad = False
                     adapter['lora_B'].requires_grad = False
-                self.adapter_scales[idx].requires_grad = True   #  ← 新增
+                    # self.adapter_scales[idx].requires_grad = True   #  ← 新增
             logging.info(f"[CurriculumLoRALinear {module_id}] Stage 0: 基础层可训练，所有 adapter 均冻结。")
             logging.info(f"[self.adapter_scales:] {self.adapter_scales}")
+            logging.info(f"self.adapter_scales[0].item(): {self.adapter_scales[0]().item()}")
+
         else:
             # 阶段大于 0，冻结基础层
             self.weight.requires_grad = False
             if self.bias is not None:
                 self.bias.requires_grad = False
             for idx, adapter in enumerate(self.adapters):
-                self.adapter_scales[idx].requires_grad = True   #  ← 新增
+                # self.adapter_scales[idx].requires_grad = True   #  ← 新增
                 logging.info(f"[self.adapter_scales:] {self.adapter_scales}")
+                logging.info(f"self.adapter_scales[0].item(): {self.adapter_scales[0]().item()}")
+
                 if idx == stage - 1:
                     adapter['lora_A'].requires_grad = True
                     adapter['lora_B'].requires_grad = True
@@ -154,9 +181,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             adapter = self.adapters[idx]
             out = F.linear(self.lora_dropout(x), adapter['lora_A'])
             out = F.linear(out, adapter['lora_B'])
-            scale = self.adapter_scales[idx] # TODO
+            scale = self.adapter_scales[idx]() # TODO: 所有adapter  对应的scale都参与训练
             if idx == self.curriculum_stage - 1:
-                adapter_out = adapter_out + self.scaling * out * scale  # 当前 adapter参与更新
+                adapter_out = adapter_out + self.scaling * out * scale  # 仅当前 adapter 参与更新
             else:
                 adapter_out = adapter_out + self.scaling * out.detach() * scale
         return baseline_out + adapter_out