fix: broadcast ref/old model state_dict for FSDP

- Add broadcast_state_dict to sync state_dict from rank-0 to all ranks
- Fix create_ref_model returning None on non-rank-0 under FSDP
- Fix sync_old_model only updating old_model on rank-0 under FSDP
- Split skip_no_cuda/skip_no_kernel markers and hoist to top-level conftest
- Add distributed tests for broadcast_state_dict and create_ref_model
This commit is contained in:
2026-07-31 08:32:22 +08:00
parent 28d1bd07cf
commit 738cb8f128
9 changed files with 248 additions and 17 deletions
+3
View File
@@ -9,6 +9,7 @@ import torch.nn.functional as F
from torch import Tensor
from astrai.factory import BaseFactory
from astrai.parallel.executor import broadcast_state_dict
from astrai.trainer.rollout import RolloutResult
@@ -391,6 +392,8 @@ class GRPOStrategy(BaseStrategy):
def sync_old_model(self):
"""Copy current policy weights to old model."""
state_dict = self.executor.unwrap_model(self.model)
if self.executor.use_distributed:
state_dict = broadcast_state_dict(state_dict)
if state_dict is not None:
self.old_model.load_state_dict(state_dict)