cleanup for failing test (#2436)

2025-03-22 17:53:29 -04:00
parent e44953d50c
commit 86bac48d14
2 changed files with 25 additions and 23 deletions
--- a/src/axolotl/monkeypatch/attention/ring_attn.py
+++ b/src/axolotl/monkeypatch/attention/ring_attn.py
@@ -27,7 +27,7 @@ def get_ring_attn_group() -> dist.ProcessGroup:
    return RING_ATTN_GROUP


-def set_ring_attn_group(ring_attn_group: dist.ProcessGroup):
+def set_ring_attn_group(ring_attn_group: dist.ProcessGroup | None):
    """
    Setter for ring attention group on this rank.