)]}'
{
  "commit": "aec09eeb3a570bb79c642e8ae540b3981e7d69cd",
  "tree": "a72ec88e0efb3d449b568da5ca597ccae9b2071d",
  "parents": [
    "e81ccfd1edfa5e753d34533900c341950402e084"
  ],
  "author": {
    "name": "Andrew Gu",
    "email": "andgu@fb.com",
    "time": "Tue Dec 20 12:56:45 2022 +0000"
  },
  "committer": {
    "name": "PyTorch MergeBot",
    "email": "pytorchmergebot@users.noreply.github.com",
    "time": "Tue Dec 20 16:49:18 2022 +0000"
  },
  "message": "[FSDP][7/N] Support `replicate` in `fully_shard` (#91044)\n\nThis PR supports nesting `replicate` in `fully_shard`.\n- The PR achieves this by treating `replicate`-annotated modules are ignored modules. This means that all submodules in the `replicate`-annotated module\u0027s subtree are ignored, including nested `fully_shard`-annotated modules, which is the desired behavior.\n\n---\n\nThis PR reworks some tree traversal.\n\nOne end goal is for `state._handles` to follow the same order for both the wrapper and composable paths. This implies that `_get_fsdp_handles()` returns the same value for both paths.\n- The helper function `_get_fully_sharded_module_to_states()` now follows a left-to-right DFS from each fully sharded module instead of a BFS. The left-to-right DFS follows `.modules()` order.\n- The composable auto \"wrap\" initialization function `_init_param_handles_from_module()` follows the reverse left-to-right DFS order. As noted in the code comments, this initialization order is a valid reverse topological sort, but it differs from the wrapper path. This is the _only_ difference with respect to initialization order through the entire process.\n```\nmod: Module(\n    submod1: Submodule()\n    submod2: Submodule(\n        subsubmod: Subsubmodule(),\n    ),\n)\n```\nFor left-to-right DFS, the order is `mod`, `submod1`, `submod2`, `subsubmod`. (For context, right-to-left DFS would be `mod`, `submod2`, `subsubmod`, `submod1`. In other words, the left-to-right vs. right-to-left corresponds to `.children()` vs. `reversed(.children())` respectively.) Then, reverse left-to-right DFS is `subsubmod`, `submod2`, `submod1`, `mod`, which is a valid initialization order. However, the wrapper auto wrap initialization order would be `submod1`, `subsubmod`, `submod2`, `mod` since it directly follows a left-to-right DFS and initializes as a part of the recursive DFS logic.\n- At the end of `_init_param_handles_from_module()`, we reverse the newly populated `state._handles`, so this is the reverse reverse left-to-right DFS order, which is equivalent to the left-to-right DFS order. Thus, `state._handles` has the same order for both paths.\n\nAnother goal is for `_get_fsdp_states()` to not traverse into any submodule that is annotated with an API that is not compatible with `fully_shard` (e.g. `replicate`). To achieve this while preserving that `_get_fsdp_states()` follows `.modules()` order, we again use a left-to-right DFS.\n\nThe reason the DFSs may look strange is because I implemented them non-recursively, which requires a stack.\n\n- `test_get_fully_sharded_module_to_states()` in `test_utils.py` checks the traversal order of `_get_fully_sharded_module_to_states()`.\n- `test_policy()` in `test_fully_shard.py` checks the traversal order returned by `_get_fsdp_handles()`.\n\n---\n\nDue to a circular dependency issue, we must move the graph/tree traversal helpers to their own file `_traversal_utils.py`, and any usages must import the entire file like `import torch.distributed.fsdp._traversal_utils as traversal_utils` instead of `from torch.distributed.fsdp._traversal_utils import ...`.\n\nThe cycle comes from the fact that the traversals require `_composable()`, which requires `_get_registry()` from `composable/contract.py`, which when imported, imports `composable/fully_shard.py`, which requires the traversals.\nPull Request resolved: https://github.com/pytorch/pytorch/pull/91044\nApproved by: https://github.com/mrshenli\n",
  "tree_diff": [
    {
      "type": "modify",
      "old_id": "3500a35f3d57595ecbfb67d121f42df2486d8f51",
      "old_mode": 33188,
      "old_path": "test/distributed/_composable/test_compose.py",
      "new_id": "643530d4a68e14ef54ce144d3ccd4879d9c3bbe7",
      "new_mode": 33188,
      "new_path": "test/distributed/_composable/test_compose.py"
    },
    {
      "type": "modify",
      "old_id": "da14aac63443aa57d5af378b66344445c4344af0",
      "old_mode": 33188,
      "old_path": "test/distributed/_composable/test_fully_shard.py",
      "new_id": "6eb59ba378c03928fbea52585043dffdaaad3fdc",
      "new_mode": 33188,
      "new_path": "test/distributed/_composable/test_fully_shard.py"
    },
    {
      "type": "modify",
      "old_id": "2f7d3e264ef0f2df855730fd2db15d9cfa759329",
      "old_mode": 33188,
      "old_path": "test/distributed/fsdp/test_fsdp_hybrid_shard.py",
      "new_id": "aa2e226c18e0e6f1dc3b6781152dfa4eca7aa2f4",
      "new_mode": 33188,
      "new_path": "test/distributed/fsdp/test_fsdp_hybrid_shard.py"
    },
    {
      "type": "modify",
      "old_id": "8bd5354b2b7015e297b950dfdcc6731c2754b732",
      "old_mode": 33188,
      "old_path": "test/distributed/fsdp/test_fsdp_overlap.py",
      "new_id": "1eb3f7ab55f740de45085a812404047cf94c9420",
      "new_mode": 33188,
      "new_path": "test/distributed/fsdp/test_fsdp_overlap.py"
    },
    {
      "type": "modify",
      "old_id": "b5ad65a9c1483ebd97df6828dd42bd1e8b563767",
      "old_mode": 33188,
      "old_path": "test/distributed/fsdp/test_utils.py",
      "new_id": "b9f96778f1da12bb4f2c587331ab9e6629a534f4",
      "new_mode": 33188,
      "new_path": "test/distributed/fsdp/test_utils.py"
    },
    {
      "type": "modify",
      "old_id": "f71b8eea859e3c0be18d4947d1214766c9c24df8",
      "old_mode": 33188,
      "old_path": "torch/distributed/_composable/fully_shard.py",
      "new_id": "5065761e0f7bfb8c01230f5e64a1b5eea8fa673b",
      "new_mode": 33188,
      "new_path": "torch/distributed/_composable/fully_shard.py"
    },
    {
      "type": "modify",
      "old_id": "322c1fe88dbe03d235884ff5f3d8675eb058e2af",
      "old_mode": 33188,
      "old_path": "torch/distributed/fsdp/_common_utils.py",
      "new_id": "52125de5bd395ead007470ae0381779ab4c021f3",
      "new_mode": 33188,
      "new_path": "torch/distributed/fsdp/_common_utils.py"
    },
    {
      "type": "modify",
      "old_id": "73b63d45366e3c9bffab1a85fd1bae48992229ea",
      "old_mode": 33188,
      "old_path": "torch/distributed/fsdp/_exec_order_utils.py",
      "new_id": "dd8a2b22878c8f93463586bd67b5d5bbe16cd1b8",
      "new_mode": 33188,
      "new_path": "torch/distributed/fsdp/_exec_order_utils.py"
    },
    {
      "type": "modify",
      "old_id": "83aea5915f44357ee0105061eb42f003fc11e7d0",
      "old_mode": 33188,
      "old_path": "torch/distributed/fsdp/_init_utils.py",
      "new_id": "2d0107d85b40448f7abfb0e93fae84ac836f1eca",
      "new_mode": 33188,
      "new_path": "torch/distributed/fsdp/_init_utils.py"
    },
    {
      "type": "modify",
      "old_id": "2b6e1912e248b1dfc754459e8c7a247b20adaa70",
      "old_mode": 33188,
      "old_path": "torch/distributed/fsdp/_optim_utils.py",
      "new_id": "6b2c635736f4fdd4276e3f65735f5f50864977e2",
      "new_mode": 33188,
      "new_path": "torch/distributed/fsdp/_optim_utils.py"
    },
    {
      "type": "modify",
      "old_id": "bac37a84cb4d5b6793234293f273a3e38e67e0dd",
      "old_mode": 33188,
      "old_path": "torch/distributed/fsdp/_runtime_utils.py",
      "new_id": "b66e4aa1087fc3e6a9246941bc6ee49ae31243e5",
      "new_mode": 33188,
      "new_path": "torch/distributed/fsdp/_runtime_utils.py"
    },
    {
      "type": "modify",
      "old_id": "ac9d440aa87a59a1346ff931176f267ee903602f",
      "old_mode": 33188,
      "old_path": "torch/distributed/fsdp/_state_dict_utils.py",
      "new_id": "81a8ea4f4f7bdd791798c82395dc390f747348c7",
      "new_mode": 33188,
      "new_path": "torch/distributed/fsdp/_state_dict_utils.py"
    },
    {
      "type": "add",
      "old_id": "0000000000000000000000000000000000000000",
      "old_mode": 0,
      "old_path": "/dev/null",
      "new_id": "86073234e542bd56174ad5920d9444b1f6d7047d",
      "new_mode": 33188,
      "new_path": "torch/distributed/fsdp/_traversal_utils.py"
    },
    {
      "type": "modify",
      "old_id": "95331874a8829818872a5b18bcf80dd7c69192d7",
      "old_mode": 33188,
      "old_path": "torch/distributed/fsdp/_wrap_utils.py",
      "new_id": "f9b5f8975486ce7dbb4eba70b65ae2c5c46bc447",
      "new_mode": 33188,
      "new_path": "torch/distributed/fsdp/_wrap_utils.py"
    },
    {
      "type": "modify",
      "old_id": "586abc4beb551b2344cfc17a942c3b6a877972ca",
      "old_mode": 33188,
      "old_path": "torch/distributed/fsdp/fully_sharded_data_parallel.py",
      "new_id": "3a98ffbeab7361b436ded7638a670b32bda95160",
      "new_mode": 33188,
      "new_path": "torch/distributed/fsdp/fully_sharded_data_parallel.py"
    },
    {
      "type": "modify",
      "old_id": "42029f153eabe01642114a25ff279d3ab980fb93",
      "old_mode": 33188,
      "old_path": "torch/testing/_internal/common_dist_composable.py",
      "new_id": "9ec92c6e71e0087662f3c0781a80cf70e9d58f58",
      "new_mode": 33188,
      "new_path": "torch/testing/_internal/common_dist_composable.py"
    }
  ]
}