Plumb use_collectives knob from KnobOptions to dcp.async_save (#1064)

Summary:
Pull Request resolved: #1064

### This Diff

Adds a `use_collectives: bool = True` field to `KnobOptions` and threads it through `DistributedCheckpointSaver._checkpoint_impl` to both `dcp.async_save(...)` and `dcp.save(...)`.

When set to `False`, DCP skips the cross-rank gather/scatter rounds it normally uses to coordinate the global save plan. Each rank writes its local shards independently. The `use_collectives` parameter is a real public kwarg on `dcp.{save,async_save}` (see `caffe2/torch/distributed/checkpoint/state_dict_saver.py`), added in upstream PyTorch PR #147758 specifically for cases where the coord-side rounds become "prohibitively expensive as the job scale grows".

### Motivation

At 397B TP=8 EP=16 (qwen35_moe_sp project), the gather/scatter rounds account for ~50 minutes of a ~67 minute save wall. Setting `use_collectives=False` removes that bottleneck entirely.

Trade-off documented in the docstring: the resulting checkpoint is no longer cross-world re-shardable; restart with the same parallelism layout still works. This is the correct default-True behavior — existing users are unaffected; the knob is opt-in.

Default `True` preserves backward compatibility for every existing torchtnt consumer.

Reviewed By: galrotem

Differential Revision: D105451194

fbshipit-source-id: f3eb80236c83793b37021c0fa023365c59cd8db4

Author: diego-urgell

torchtnt/framework/callbacks/checkpointer_types.py

@@ -20,6 +20,13 @@ class KnobOptions:
         max_per_rank_io_concurrency: Maximum number of concurrent IO operations per rank in checkpointing.
                                      Defaults to 16.
         enable_storage_optimization: Enable storage efficiency optimizations for Distributed Checkpointing.
+        use_collectives: If ``False``, skip the cross-rank gather/scatter rounds that DCP uses to
+                         coordinate the global save plan. Each rank writes its local shards
+                         independently. Defaults to ``True``. Setting to ``False`` removes the
+                         coord-side bottleneck that becomes prohibitively expensive at large scales
+                         (e.g. ~50 minutes of a ~67 minute save wall at 397B TP=8 EP=16). Trade-off:
+                         the resulting checkpoint is no longer cross-world re-shardable; restart
+                         with the same parallelism layout still works.
     """
 
     # use a more conservative number of concurrent IO operations per rank in Checkpointing
@@ -28,6 +35,8 @@ class KnobOptions:
     # This would enable storage efficiency optimizations (model store):
     # e.g. Compression, Batching, Quantization etc.
     enable_storage_optimization: bool = True
+    # Skip DCP coord-side gather/scatter rounds.
+    use_collectives: bool = True
 
 
 @dataclass

torchtnt/framework/callbacks/dcp_saver.py

@@ -198,6 +198,7 @@ def _checkpoint_impl(
                     storage_writer=storage_writer,
                     planner=planner,
                     async_stager=stager,
+                    use_collectives=self._knob_options.use_collectives,
                 )
                 self._prev_snapshot = cast(Future, prev_snapshot)
                 if curr_snapshot_wait:
@@ -210,6 +211,7 @@ def _checkpoint_impl(
                     process_group=self._process_group,
                     storage_writer=storage_writer,
                     planner=planner,
+                    use_collectives=self._knob_options.use_collectives,
                 )
 
         return True