fix comments

Azure · Jan 11, 2024 · d3acfa2 · d3acfa2
1 parent a8ff215
commit d3acfa2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/msamp/fsdp/_runtime_utils.py b/msamp/fsdp/_runtime_utils.py
@@ -15,8 +15,8 @@
 def _fp8_post_backward_hook(state, handle, *unused):
     """A post-backward communication hook which supports fp8."""
     accumulate_grad = hasattr(state._flat_param, '_saved_grad_shard')
-    if accumulate_grad and not torch.all(state._flat_param._saved_grad_shard == 0):
-        raise NotImplementedError('accumulate_grad is not supported for fp8')
+    if accumulate_grad and torch.count_nonzero(state._flat_param._saved_grad_shard).item() > 0:
+        raise NotImplementedError('accumulate_grad is not supported yet for fp8')
 
     old_communication_hook = state._communication_hook
     state._communication_hook = state._get_fp8_comm_hook()