Merge pull request baidu-baige#158 from NeosZhang/feat/ddp-comm-hook-v2

[train] feat: add ddp_comm_hook for gradient compression

Author: nullnonenilNULL

loongforge/embodied/distributed/ddp_utils/__init__.py

@@ -0,0 +1,8 @@
+# Copyright 2026 The LoongForge Authors.
+# SPDX-License-Identifier: Apache-2.0
+
+"""DDP utilities package: gradient communication hook resolution."""
+
+from .ddp_comm_hook import resolve_comm_hook
+
+__all__ = ["resolve_comm_hook"]

loongforge/embodied/distributed/ddp_utils/ddp_comm_hook.py

@@ -0,0 +1,73 @@
+# Copyright 2026 The LoongForge Authors.
+# SPDX-License-Identifier: Apache-2.0
+
+"""DDP gradient communication hooks."""
+
+import logging
+from typing import Callable
+
+from torch.distributed.algorithms.ddp_comm_hooks.default_hooks import (
+    allreduce_hook,
+    bf16_compress_hook,
+    fp16_compress_hook,
+)
+
+from ..utils import is_rank_zero
+
+logger = logging.getLogger(__name__)
+
+# Comm hooks eligible to be resolved by ``resolve_comm_hook``. All of them
+# share the ``(process_group, bucket) -> Future`` signature.
+_SUPPORTED_COMM_HOOKS = {
+    "allreduce_hook": allreduce_hook,
+    "fp16_compress_hook": fp16_compress_hook,
+    "bf16_compress_hook": bf16_compress_hook,
+}
+
+
+def resolve_comm_hook(hook: Callable | str, use_logging: bool = False) -> Callable:
+    """Resolve a comm hook name/callable, optionally wrapped with logging.
+
+    Args:
+        hook: Either a DDP comm hook callable (e.g. ``allreduce_hook``,
+            ``fp16_compress_hook``, ``bf16_compress_hook``) or its name.
+        use_logging: If True, wrap the resolved hook so it logs bucket info
+            on rank 0 before/after it runs.
+
+    Returns:
+        A comm hook with the ``(process_group, bucket)`` signature.
+    """
+    if isinstance(hook, str):
+        try:
+            hook = _SUPPORTED_COMM_HOOKS[hook]
+        except KeyError:
+            raise ValueError(
+                f"Unsupported comm hook name {hook!r}, expected one of "
+                f"{sorted(_SUPPORTED_COMM_HOOKS)}."
+            ) from None
+
+    if not use_logging:
+        return hook
+
+    hook_name = hook.__name__
+
+    def logging_comm_hook(process_group, bucket):
+        if is_rank_zero():
+            tensor = bucket.buffer()
+            logger.info(
+                "DDP %s: bucket_index=%d numel=%d dtype=%s",
+                hook_name,
+                bucket.index(),
+                tensor.numel(),
+                tensor.dtype,
+            )
+        fut = hook(process_group, bucket)
+        if is_rank_zero():
+            fut.add_done_callback(
+                lambda fut: logger.info(
+                    "DDP %s done: bucket_index=%d", hook_name, bucket.index()
+                )
+            )
+        return fut
+
+    return logging_comm_hook

loongforge/embodied/distributed/parallel.py

@@ -14,6 +14,7 @@
 
 from .activation_checkpointing import apply_activation_checkpointing
 from .context import DistributedContext
+from .ddp_utils import resolve_comm_hook
 from .utils import (
     filter_supported_kwargs,
     get_module_names_by_dtype,
@@ -93,7 +94,16 @@ def _wrap_ddp(model: nn.Module, training_args, ctx: DistributedContext, dtype: t
         "batched_grad_copy": training_args.ddp_batched_grad_copy,
     }
 
-    return DDP(model, **filter_supported_kwargs(DDP, ddp_kwargs))
+    ddp_model = DDP(model, **filter_supported_kwargs(DDP, ddp_kwargs))
+    if training_args.ddp_comm_hook:
+        # allreduce_hook: full-precision gradient all-reduce (default, no compression);
+        # fp16_compress_hook: compresses gradients to fp16 before all-reduce to cut traffic;
+        # bf16_compress_hook: compresses gradients to bf16 before all-reduce to cut traffic.
+        comm_hook = resolve_comm_hook(
+            training_args.ddp_comm_hook, use_logging=training_args.ddp_comm_hook_logging
+        )
+        ddp_model.register_comm_hook(state=None, hook=comm_hook)
+    return ddp_model
 
 
 def _wrap_fsdp(

loongforge/embodied/train/training_args.py

@@ -1164,6 +1164,20 @@ class _DistributedArgs:
             "help": "Batch gradient copies into buckets to reduce kernel launches."
         },
     )
+    ddp_comm_hook: Optional[str] = field(
+        default=None,
+        metadata={
+            "choices": ["allreduce_hook", "fp16_compress_hook", "bf16_compress_hook"],
+            "help": "DDP gradient communication hook.",
+        },
+    )
+    ddp_comm_hook_logging: bool = field(
+        default=False,
+        metadata={
+            "help": "Wrap the DDP comm hook with rank-0 logging of bucket info "
+                    "before/after each all-reduce."
+        },
+    )
     dynamo_optimize_ddp: bool = field(
         default=True,
         metadata={