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

tests/extension/conftest.py

@@ -4,13 +4,8 @@ import pytest
 import torch
 
 from astrai.config.model_config import AutoRegressiveLMConfig
-from astrai.extension import is_available
 from astrai.model.transformer import AutoRegressiveLM
-
-CUDA_AVAILABLE = torch.cuda.is_available() and is_available("attn_paged_decode")
-skip_no_cuda = pytest.mark.skipif(
-    not CUDA_AVAILABLE, reason="CUDA not available or kernels not built"
-)
+from tests.conftest import skip_no_kernel
 
 D = 64
 CFG = dict(