We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
shard_model does not respect shard_count.
## 🐛 Bug

`autoshard.shard_model()` does not respect `shard_count`. See [comments](https://github.com/facebookresearch/fairscale/pull/817/files#r733237403) in https://github.com/facebookresearch/fairscale/pull/817 for more context and how to reproduce. 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `fairscale/experimental/nn/auto_shard.py`
Content:
```
1 # Copyright (c) Facebook, Inc. and its affiliates.
2 #
3 # This source code is licensed under the BSD license found in the
4 # LICENSE file in the root directory of this source tree.
5 
6 import logging
7 from typing import Dict, List, Set
8 
9 import torch
10 import torch.fx
11 from torch.fx.node import Node
12 
13 
14 def _get_count(param_count: Dict, node_name: str) -> int:
15     """Identify different mutations of a given node name."""
16     # TODO(anj): This is not very stable since it is possible that the name
17     # may not be in the same format. Is there another way to identify nodes
18     # in a graph?
19     if node_name in param_count:
20         return param_count[node_name]
21     elif node_name.split("_")[0] in param_count:
22         return param_count[node_name.split("_")[0]]
23     else:
24         raise RuntimeError(f"Unable to find match between param {param_count} and node {node_name}")
25 
26 
27 def _create_shard_to_param_count(param_count: Dict, node_name_to_shard_id: Dict) -> Dict:
28     """Utility to create a map from shard id to param count using existing state."""
29 
30     shard_to_param_count: Dict[int, int] = {}
31     for node_name in node_name_to_shard_id.keys():
32         try:
33             count = _get_count(param_count, node_name)
34         except RuntimeError:
35             continue
36         if node_name_to_shard_id[node_name] in shard_to_param_count:
37             shard_to_param_count[node_name_to_shard_id[node_name]] += count
38         else:
39             shard_to_param_count[node_name_to_shard_id[node_name]] = count
40     return shard_to_param_count
41 
42 
43 def _split_nodes(traced_graph_module: torch.fx.GraphModule, shard_count: int = 3) -> Dict:
44     """Utility used to trace a graph and identify shard cutpoints."""
45 
46     node_name_to_shard_id: Dict[str, int] = {}
47     shard_id = 0
48     nodes_so_far = []
49     param_count: Dict[str, int] = {}
50     shard_to_param_count = {}
51 
52     # Find the total number of params in the model and
53     # the number of params per shard we are aiming for.
54     for name, module in traced_graph_module.named_modules():
55         if "." in name:
56             continue
57         param_count[name] = sum([x.numel() for x in module.parameters()])
58     logging.info(f"Total number of params are {param_count['']}")
59     per_shard_param = param_count[""] // shard_count
60     logging.info(f"Per shard param count {per_shard_param}")
61 
62     for node in traced_graph_module.graph.nodes:
63         if node.op == "placeholder":
64             node_name_to_shard_id[node.name] = shard_id
65             nodes_so_far.append(node.name)
66         elif node.op in ["get_attr", "call_function", "call_method", "call_module"]:
67 
68             min_shard_id = shard_id
69             min_node_name = ""
70             # For each of the args of a given node, find the arg that is not the
71             # last node we traversed. This is to help us find skip connections
72             # across shards.
73             for arg in node.args:
74                 # If the node has args that are inputs to the forward function, they
75                 # may not have explicit names.
76                 if not hasattr(arg, "name"):
77                     continue
78 
79                 if arg.name in node_name_to_shard_id and arg.name != nodes_so_far[-1]:
80                     if node_name_to_shard_id[arg.name] < min_shard_id:
81                         min_shard_id = node_name_to_shard_id[arg.name]
82                         min_node_name = arg.name
83 
84             # If there is an input that is not from the previous shard,
85             # we collapse all the shards in between to be part of 1 shard.
86             # and update the param count per shard accordingly.
87             if min_shard_id < shard_id:
88                 for node_name in reversed(nodes_so_far):
89                     node_name_to_shard_id[node_name] = min_shard_id
90                     if node_name == min_node_name:
91                         break
92                 shard_id = min_shard_id
93                 # TODO(anj-s): Find a way to raise an error early if this can cause OOM errors.
94                 shard_to_param_count = _create_shard_to_param_count(param_count, node_name_to_shard_id)
95 
96             # Update state that is tracking node -> shard id and shard id -> param count.
97             node_name_to_shard_id[node.name] = shard_id
98             nodes_so_far.append(node.name)
99             # TODO(anj): This could just be an update, we don't need to recreate the map.
100             shard_to_param_count = _create_shard_to_param_count(param_count, node_name_to_shard_id)
101             # If we have gone over the number of params per shard count that we want to
102             # achieve, we should add a new shard.
103             # The shard_id may not have been updated in the map if we are at a node that does not
104             # have params.
105             if shard_id in shard_to_param_count and shard_to_param_count[shard_id] > per_shard_param:
106                 shard_id += 1
107         elif node.op == "output":
108             break
109     return node_name_to_shard_id
110 
111 
112 class _ExtendedLeafTracer(torch.fx.Tracer):
113     """Tracer with an extended set of leaf nn.Modules."""
114 
115     def __init__(self, leaf_modules: Set[torch.nn.Module]):
116         """Initializes a new _ExtendedLeafTracer object.
117 
118         Args:
119             leaf_modules: The set of extra nn.Modules instances which will not be traced
120                 through but instead considered to be leaves.
121         """
122         super().__init__()
123         self.leaf_modules = leaf_modules
124 
125     def is_leaf_module(self, m: torch.nn.Module, model_qualified_name: str) -> bool:
126         return super().is_leaf_module(m, model_qualified_name) or m in self.leaf_modules
127 
128 
129 # TODO(ehotaj): Extend this method to wrap at the least granular level. One way to do
130 # would be to wrap the Module tree bottom up, first wrapping untracable children and
131 # only wrapping parents if they are also untracable.
132 def _trace(model: torch.nn.Module) -> torch.fx.GraphModule:
133     """Traces the given model and automatically wraps untracable modules into leaves."""
134     leaf_modules = set()
135     tracer = _ExtendedLeafTracer(leaf_modules)
136     for name, module in model.named_modules():
137         # TODO(ehotaj): The default is_leaf_module includes everything in torch.nn.
138         # This means that some coarse modules like nn.TransformerEncoder are treated
139         # as leaves, not traced, and are unable to be sharded. We may want to extend our
140         # sharding code to trace through these modules as well.
141         if tracer.is_leaf_module(module, ""):
142             continue
143         try:
144             tracer.trace(module)
145         except (TypeError, torch.fx.proxy.TraceError):
146             leaf_modules.add(module)
147             tracer = _ExtendedLeafTracer(leaf_modules)
148     graph = tracer.trace(model)
149     return torch.fx.GraphModule(model, graph)
150 
151 
152 def shard_model(model: torch.nn.Module, shard_count: int = 3) -> List[torch.fx.GraphModule]:
153     """Utility used to shard a model using torch.fx.
154 
155     This function traces the model twice in an attempt to identify the
156     right cutpoints and then shard the model. In the first pass we calculate
157     the number of parameters as we are tracing the graph and mark nodes at
158     which we might want to create a new module. In the second pass we
159     modify the graph by inserting placeholders and output nodes to essentially
160     shard the graph.
161 
162     We don't support skip connections between shards. This means that all
163     input and output is self contained within a given shard. A node from
164     shard 1 cannot be an input to a node from shard 3. We expect all inputs
165     to a given shard to be coming from the last node in the previous shard.
166     This means that we may not be able to shard models by the specified
167     `shard_count` mentioned by the user.
168 
169     Args:
170         model (nn.Module): Model to be sharded as specified by the device count.
171 
172         shard_count (int): Number of shards that we want to split the model into.
173 
174     """
175     module_list: List[torch.fx.GraphModule] = []
176     num_graphs = 0
177     new_graph = torch.fx.Graph()  # type: ignore
178     env: Dict[str, Node] = {}
179     new_input_node = None
180 
181     traced_graph_module = _trace(model)
182 
183     # This is the first pass where we attempt to get a map of where
184     # we need to insert placeholder and output nodes.
185     node_name_to_shard_id = _split_nodes(traced_graph_module, shard_count=shard_count)
186 
187     # dummy value which indicates that this is the first node.
188     prev_shard_id = 1000
189     prev_node = None
190     for node in traced_graph_module.graph.nodes:
191         # If the current node is in the next shard, we insert an output node.
192         # A new graph is created and a placeholder is added for the next shard.
193         if node.name in node_name_to_shard_id and prev_shard_id < node_name_to_shard_id[node.name]:
194             assert prev_node, "prev_node cannot be None"
195 
196             with new_graph.inserting_after(prev_node):
197                 new_graph.output(env[prev_node.name])
198             num_graphs += 1
199             module_list.append(torch.fx.GraphModule(model, new_graph))
200             new_graph = torch.fx.Graph()
201             node_name = "placeholder" + str(num_graphs)
202             pl_node = new_graph.create_node("placeholder", node_name)
203             env[node_name] = pl_node
204             new_input_node = pl_node
205 
206         if new_input_node is not None:
207             # Account for a placeholder in the new graph.
208             node.args = (new_input_node,)
209             new_input_node = None
210         if node.op in ["placeholder", "get_attr", "call_function", "call_method", "call_module"]:
211             # Copy the nodes from the existing graph to the new graph.
212             new_node = new_graph.node_copy(node, lambda x: env[x.name])
213             env[node.name] = new_node
214         elif node.op == "output":
215             # If this is the last node, we should add an output
216             # node and add the last graph to the list.
217             assert prev_node, "prev_node cannot be None"
218 
219             with new_graph.inserting_after(prev_node):
220                 new_graph.output(env[prev_node.name])
221             module_list.append(torch.fx.GraphModule(model, new_graph))
222             break
223         prev_node = new_node
224         prev_shard_id = node_name_to_shard_id[node.name]
225 
226     return module_list
227
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/fairscale/experimental/nn/auto_shard.py b/fairscale/experimental/nn/auto_shard.py
--- a/fairscale/experimental/nn/auto_shard.py
+++ b/fairscale/experimental/nn/auto_shard.py
@@ -52,8 +52,7 @@
     # Find the total number of params in the model and
     # the number of params per shard we are aiming for.
     for name, module in traced_graph_module.named_modules():
-        if "." in name:
-            continue
+        name = name.replace(".", "_")
         param_count[name] = sum([x.numel() for x in module.parameters()])
     logging.info(f"Total number of params are {param_count['']}")
     per_shard_param = param_count[""] // shard_count

{"golden_diff": "diff --git a/fairscale/experimental/nn/auto_shard.py b/fairscale/experimental/nn/auto_shard.py\n--- a/fairscale/experimental/nn/auto_shard.py\n+++ b/fairscale/experimental/nn/auto_shard.py\n@@ -52,8 +52,7 @@\n     # Find the total number of params in the model and\n     # the number of params per shard we are aiming for.\n     for name, module in traced_graph_module.named_modules():\n-        if \".\" in name:\n-            continue\n+        name = name.replace(\".\", \"_\")\n         param_count[name] = sum([x.numel() for x in module.parameters()])\n     logging.info(f\"Total number of params are {param_count['']}\")\n     per_shard_param = param_count[\"\"] // shard_count\n", "issue": "shard_model does not respect shard_count.\n## \ud83d\udc1b Bug\r\n\r\n`autoshard.shard_model()` does not respect `shard_count`. See [comments](https://github.com/facebookresearch/fairscale/pull/817/files#r733237403) in https://github.com/facebookresearch/fairscale/pull/817 for more context and how to reproduce. \n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\nfrom typing import Dict, List, Set\n\nimport torch\nimport torch.fx\nfrom torch.fx.node import Node\n\n\ndef _get_count(param_count: Dict, node_name: str) -> int:\n    \"\"\"Identify different mutations of a given node name.\"\"\"\n    # TODO(anj): This is not very stable since it is possible that the name\n    # may not be in the same format. Is there another way to identify nodes\n    # in a graph?\n    if node_name in param_count:\n        return param_count[node_name]\n    elif node_name.split(\"_\")[0] in param_count:\n        return param_count[node_name.split(\"_\")[0]]\n    else:\n        raise RuntimeError(f\"Unable to find match between param {param_count} and node {node_name}\")\n\n\ndef _create_shard_to_param_count(param_count: Dict, node_name_to_shard_id: Dict) -> Dict:\n    \"\"\"Utility to create a map from shard id to param count using existing state.\"\"\"\n\n    shard_to_param_count: Dict[int, int] = {}\n    for node_name in node_name_to_shard_id.keys():\n        try:\n            count = _get_count(param_count, node_name)\n        except RuntimeError:\n            continue\n        if node_name_to_shard_id[node_name] in shard_to_param_count:\n            shard_to_param_count[node_name_to_shard_id[node_name]] += count\n        else:\n            shard_to_param_count[node_name_to_shard_id[node_name]] = count\n    return shard_to_param_count\n\n\ndef _split_nodes(traced_graph_module: torch.fx.GraphModule, shard_count: int = 3) -> Dict:\n    \"\"\"Utility used to trace a graph and identify shard cutpoints.\"\"\"\n\n    node_name_to_shard_id: Dict[str, int] = {}\n    shard_id = 0\n    nodes_so_far = []\n    param_count: Dict[str, int] = {}\n    shard_to_param_count = {}\n\n    # Find the total number of params in the model and\n    # the number of params per shard we are aiming for.\n    for name, module in traced_graph_module.named_modules():\n        if \".\" in name:\n            continue\n        param_count[name] = sum([x.numel() for x in module.parameters()])\n    logging.info(f\"Total number of params are {param_count['']}\")\n    per_shard_param = param_count[\"\"] // shard_count\n    logging.info(f\"Per shard param count {per_shard_param}\")\n\n    for node in traced_graph_module.graph.nodes:\n        if node.op == \"placeholder\":\n            node_name_to_shard_id[node.name] = shard_id\n            nodes_so_far.append(node.name)\n        elif node.op in [\"get_attr\", \"call_function\", \"call_method\", \"call_module\"]:\n\n            min_shard_id = shard_id\n            min_node_name = \"\"\n            # For each of the args of a given node, find the arg that is not the\n            # last node we traversed. This is to help us find skip connections\n            # across shards.\n            for arg in node.args:\n                # If the node has args that are inputs to the forward function, they\n                # may not have explicit names.\n                if not hasattr(arg, \"name\"):\n                    continue\n\n                if arg.name in node_name_to_shard_id and arg.name != nodes_so_far[-1]:\n                    if node_name_to_shard_id[arg.name] < min_shard_id:\n                        min_shard_id = node_name_to_shard_id[arg.name]\n                        min_node_name = arg.name\n\n            # If there is an input that is not from the previous shard,\n            # we collapse all the shards in between to be part of 1 shard.\n            # and update the param count per shard accordingly.\n            if min_shard_id < shard_id:\n                for node_name in reversed(nodes_so_far):\n                    node_name_to_shard_id[node_name] = min_shard_id\n                    if node_name == min_node_name:\n                        break\n                shard_id = min_shard_id\n                # TODO(anj-s): Find a way to raise an error early if this can cause OOM errors.\n                shard_to_param_count = _create_shard_to_param_count(param_count, node_name_to_shard_id)\n\n            # Update state that is tracking node -> shard id and shard id -> param count.\n            node_name_to_shard_id[node.name] = shard_id\n            nodes_so_far.append(node.name)\n            # TODO(anj): This could just be an update, we don't need to recreate the map.\n            shard_to_param_count = _create_shard_to_param_count(param_count, node_name_to_shard_id)\n            # If we have gone over the number of params per shard count that we want to\n            # achieve, we should add a new shard.\n            # The shard_id may not have been updated in the map if we are at a node that does not\n            # have params.\n            if shard_id in shard_to_param_count and shard_to_param_count[shard_id] > per_shard_param:\n                shard_id += 1\n        elif node.op == \"output\":\n            break\n    return node_name_to_shard_id\n\n\nclass _ExtendedLeafTracer(torch.fx.Tracer):\n    \"\"\"Tracer with an extended set of leaf nn.Modules.\"\"\"\n\n    def __init__(self, leaf_modules: Set[torch.nn.Module]):\n        \"\"\"Initializes a new _ExtendedLeafTracer object.\n\n        Args:\n            leaf_modules: The set of extra nn.Modules instances which will not be traced\n                through but instead considered to be leaves.\n        \"\"\"\n        super().__init__()\n        self.leaf_modules = leaf_modules\n\n    def is_leaf_module(self, m: torch.nn.Module, model_qualified_name: str) -> bool:\n        return super().is_leaf_module(m, model_qualified_name) or m in self.leaf_modules\n\n\n# TODO(ehotaj): Extend this method to wrap at the least granular level. One way to do\n# would be to wrap the Module tree bottom up, first wrapping untracable children and\n# only wrapping parents if they are also untracable.\ndef _trace(model: torch.nn.Module) -> torch.fx.GraphModule:\n    \"\"\"Traces the given model and automatically wraps untracable modules into leaves.\"\"\"\n    leaf_modules = set()\n    tracer = _ExtendedLeafTracer(leaf_modules)\n    for name, module in model.named_modules():\n        # TODO(ehotaj): The default is_leaf_module includes everything in torch.nn.\n        # This means that some coarse modules like nn.TransformerEncoder are treated\n        # as leaves, not traced, and are unable to be sharded. We may want to extend our\n        # sharding code to trace through these modules as well.\n        if tracer.is_leaf_module(module, \"\"):\n            continue\n        try:\n            tracer.trace(module)\n        except (TypeError, torch.fx.proxy.TraceError):\n            leaf_modules.add(module)\n            tracer = _ExtendedLeafTracer(leaf_modules)\n    graph = tracer.trace(model)\n    return torch.fx.GraphModule(model, graph)\n\n\ndef shard_model(model: torch.nn.Module, shard_count: int = 3) -> List[torch.fx.GraphModule]:\n    \"\"\"Utility used to shard a model using torch.fx.\n\n    This function traces the model twice in an attempt to identify the\n    right cutpoints and then shard the model. In the first pass we calculate\n    the number of parameters as we are tracing the graph and mark nodes at\n    which we might want to create a new module. In the second pass we\n    modify the graph by inserting placeholders and output nodes to essentially\n    shard the graph.\n\n    We don't support skip connections between shards. This means that all\n    input and output is self contained within a given shard. A node from\n    shard 1 cannot be an input to a node from shard 3. We expect all inputs\n    to a given shard to be coming from the last node in the previous shard.\n    This means that we may not be able to shard models by the specified\n    `shard_count` mentioned by the user.\n\n    Args:\n        model (nn.Module): Model to be sharded as specified by the device count.\n\n        shard_count (int): Number of shards that we want to split the model into.\n\n    \"\"\"\n    module_list: List[torch.fx.GraphModule] = []\n    num_graphs = 0\n    new_graph = torch.fx.Graph()  # type: ignore\n    env: Dict[str, Node] = {}\n    new_input_node = None\n\n    traced_graph_module = _trace(model)\n\n    # This is the first pass where we attempt to get a map of where\n    # we need to insert placeholder and output nodes.\n    node_name_to_shard_id = _split_nodes(traced_graph_module, shard_count=shard_count)\n\n    # dummy value which indicates that this is the first node.\n    prev_shard_id = 1000\n    prev_node = None\n    for node in traced_graph_module.graph.nodes:\n        # If the current node is in the next shard, we insert an output node.\n        # A new graph is created and a placeholder is added for the next shard.\n        if node.name in node_name_to_shard_id and prev_shard_id < node_name_to_shard_id[node.name]:\n            assert prev_node, \"prev_node cannot be None\"\n\n            with new_graph.inserting_after(prev_node):\n                new_graph.output(env[prev_node.name])\n            num_graphs += 1\n            module_list.append(torch.fx.GraphModule(model, new_graph))\n            new_graph = torch.fx.Graph()\n            node_name = \"placeholder\" + str(num_graphs)\n            pl_node = new_graph.create_node(\"placeholder\", node_name)\n            env[node_name] = pl_node\n            new_input_node = pl_node\n\n        if new_input_node is not None:\n            # Account for a placeholder in the new graph.\n            node.args = (new_input_node,)\n            new_input_node = None\n        if node.op in [\"placeholder\", \"get_attr\", \"call_function\", \"call_method\", \"call_module\"]:\n            # Copy the nodes from the existing graph to the new graph.\n            new_node = new_graph.node_copy(node, lambda x: env[x.name])\n            env[node.name] = new_node\n        elif node.op == \"output\":\n            # If this is the last node, we should add an output\n            # node and add the last graph to the list.\n            assert prev_node, \"prev_node cannot be None\"\n\n            with new_graph.inserting_after(prev_node):\n                new_graph.output(env[prev_node.name])\n            module_list.append(torch.fx.GraphModule(model, new_graph))\n            break\n        prev_node = new_node\n        prev_shard_id = node_name_to_shard_id[node.name]\n\n    return module_list\n", "path": "fairscale/experimental/nn/auto_shard.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\nfrom typing import Dict, List, Set\n\nimport torch\nimport torch.fx\nfrom torch.fx.node import Node\n\n\ndef _get_count(param_count: Dict, node_name: str) -> int:\n    \"\"\"Identify different mutations of a given node name.\"\"\"\n    # TODO(anj): This is not very stable since it is possible that the name\n    # may not be in the same format. Is there another way to identify nodes\n    # in a graph?\n    if node_name in param_count:\n        return param_count[node_name]\n    elif node_name.split(\"_\")[0] in param_count:\n        return param_count[node_name.split(\"_\")[0]]\n    else:\n        raise RuntimeError(f\"Unable to find match between param {param_count} and node {node_name}\")\n\n\ndef _create_shard_to_param_count(param_count: Dict, node_name_to_shard_id: Dict) -> Dict:\n    \"\"\"Utility to create a map from shard id to param count using existing state.\"\"\"\n\n    shard_to_param_count: Dict[int, int] = {}\n    for node_name in node_name_to_shard_id.keys():\n        try:\n            count = _get_count(param_count, node_name)\n        except RuntimeError:\n            continue\n        if node_name_to_shard_id[node_name] in shard_to_param_count:\n            shard_to_param_count[node_name_to_shard_id[node_name]] += count\n        else:\n            shard_to_param_count[node_name_to_shard_id[node_name]] = count\n    return shard_to_param_count\n\n\ndef _split_nodes(traced_graph_module: torch.fx.GraphModule, shard_count: int = 3) -> Dict:\n    \"\"\"Utility used to trace a graph and identify shard cutpoints.\"\"\"\n\n    node_name_to_shard_id: Dict[str, int] = {}\n    shard_id = 0\n    nodes_so_far = []\n    param_count: Dict[str, int] = {}\n    shard_to_param_count = {}\n\n    # Find the total number of params in the model and\n    # the number of params per shard we are aiming for.\n    for name, module in traced_graph_module.named_modules():\n        name = name.replace(\".\", \"_\")\n        param_count[name] = sum([x.numel() for x in module.parameters()])\n    logging.info(f\"Total number of params are {param_count['']}\")\n    per_shard_param = param_count[\"\"] // shard_count\n    logging.info(f\"Per shard param count {per_shard_param}\")\n\n    for node in traced_graph_module.graph.nodes:\n        if node.op == \"placeholder\":\n            node_name_to_shard_id[node.name] = shard_id\n            nodes_so_far.append(node.name)\n        elif node.op in [\"get_attr\", \"call_function\", \"call_method\", \"call_module\"]:\n\n            min_shard_id = shard_id\n            min_node_name = \"\"\n            # For each of the args of a given node, find the arg that is not the\n            # last node we traversed. This is to help us find skip connections\n            # across shards.\n            for arg in node.args:\n                # If the node has args that are inputs to the forward function, they\n                # may not have explicit names.\n                if not hasattr(arg, \"name\"):\n                    continue\n\n                if arg.name in node_name_to_shard_id and arg.name != nodes_so_far[-1]:\n                    if node_name_to_shard_id[arg.name] < min_shard_id:\n                        min_shard_id = node_name_to_shard_id[arg.name]\n                        min_node_name = arg.name\n\n            # If there is an input that is not from the previous shard,\n            # we collapse all the shards in between to be part of 1 shard.\n            # and update the param count per shard accordingly.\n            if min_shard_id < shard_id:\n                for node_name in reversed(nodes_so_far):\n                    node_name_to_shard_id[node_name] = min_shard_id\n                    if node_name == min_node_name:\n                        break\n                shard_id = min_shard_id\n                # TODO(anj-s): Find a way to raise an error early if this can cause OOM errors.\n                shard_to_param_count = _create_shard_to_param_count(param_count, node_name_to_shard_id)\n\n            # Update state that is tracking node -> shard id and shard id -> param count.\n            node_name_to_shard_id[node.name] = shard_id\n            nodes_so_far.append(node.name)\n            # TODO(anj): This could just be an update, we don't need to recreate the map.\n            shard_to_param_count = _create_shard_to_param_count(param_count, node_name_to_shard_id)\n            # If we have gone over the number of params per shard count that we want to\n            # achieve, we should add a new shard.\n            # The shard_id may not have been updated in the map if we are at a node that does not\n            # have params.\n            if shard_id in shard_to_param_count and shard_to_param_count[shard_id] > per_shard_param:\n                shard_id += 1\n        elif node.op == \"output\":\n            break\n    return node_name_to_shard_id\n\n\nclass _ExtendedLeafTracer(torch.fx.Tracer):\n    \"\"\"Tracer with an extended set of leaf nn.Modules.\"\"\"\n\n    def __init__(self, leaf_modules: Set[torch.nn.Module]):\n        \"\"\"Initializes a new _ExtendedLeafTracer object.\n\n        Args:\n            leaf_modules: The set of extra nn.Modules instances which will not be traced\n                through but instead considered to be leaves.\n        \"\"\"\n        super().__init__()\n        self.leaf_modules = leaf_modules\n\n    def is_leaf_module(self, m: torch.nn.Module, model_qualified_name: str) -> bool:\n        return super().is_leaf_module(m, model_qualified_name) or m in self.leaf_modules\n\n\n# TODO(ehotaj): Extend this method to wrap at the least granular level. One way to do\n# would be to wrap the Module tree bottom up, first wrapping untracable children and\n# only wrapping parents if they are also untracable.\ndef _trace(model: torch.nn.Module) -> torch.fx.GraphModule:\n    \"\"\"Traces the given model and automatically wraps untracable modules into leaves.\"\"\"\n    leaf_modules = set()\n    tracer = _ExtendedLeafTracer(leaf_modules)\n    for name, module in model.named_modules():\n        # TODO(ehotaj): The default is_leaf_module includes everything in torch.nn.\n        # This means that some coarse modules like nn.TransformerEncoder are treated\n        # as leaves, not traced, and are unable to be sharded. We may want to extend our\n        # sharding code to trace through these modules as well.\n        if tracer.is_leaf_module(module, \"\"):\n            continue\n        try:\n            tracer.trace(module)\n        except (TypeError, torch.fx.proxy.TraceError):\n            leaf_modules.add(module)\n            tracer = _ExtendedLeafTracer(leaf_modules)\n    graph = tracer.trace(model)\n    return torch.fx.GraphModule(model, graph)\n\n\ndef shard_model(model: torch.nn.Module, shard_count: int = 3) -> List[torch.fx.GraphModule]:\n    \"\"\"Utility used to shard a model using torch.fx.\n\n    This function traces the model twice in an attempt to identify the\n    right cutpoints and then shard the model. In the first pass we calculate\n    the number of parameters as we are tracing the graph and mark nodes at\n    which we might want to create a new module. In the second pass we\n    modify the graph by inserting placeholders and output nodes to essentially\n    shard the graph.\n\n    We don't support skip connections between shards. This means that all\n    input and output is self contained within a given shard. A node from\n    shard 1 cannot be an input to a node from shard 3. We expect all inputs\n    to a given shard to be coming from the last node in the previous shard.\n    This means that we may not be able to shard models by the specified\n    `shard_count` mentioned by the user.\n\n    Args:\n        model (nn.Module): Model to be sharded as specified by the device count.\n\n        shard_count (int): Number of shards that we want to split the model into.\n\n    \"\"\"\n    module_list: List[torch.fx.GraphModule] = []\n    num_graphs = 0\n    new_graph = torch.fx.Graph()  # type: ignore\n    env: Dict[str, Node] = {}\n    new_input_node = None\n\n    traced_graph_module = _trace(model)\n\n    # This is the first pass where we attempt to get a map of where\n    # we need to insert placeholder and output nodes.\n    node_name_to_shard_id = _split_nodes(traced_graph_module, shard_count=shard_count)\n\n    # dummy value which indicates that this is the first node.\n    prev_shard_id = 1000\n    prev_node = None\n    for node in traced_graph_module.graph.nodes:\n        # If the current node is in the next shard, we insert an output node.\n        # A new graph is created and a placeholder is added for the next shard.\n        if node.name in node_name_to_shard_id and prev_shard_id < node_name_to_shard_id[node.name]:\n            assert prev_node, \"prev_node cannot be None\"\n\n            with new_graph.inserting_after(prev_node):\n                new_graph.output(env[prev_node.name])\n            num_graphs += 1\n            module_list.append(torch.fx.GraphModule(model, new_graph))\n            new_graph = torch.fx.Graph()\n            node_name = \"placeholder\" + str(num_graphs)\n            pl_node = new_graph.create_node(\"placeholder\", node_name)\n            env[node_name] = pl_node\n            new_input_node = pl_node\n\n        if new_input_node is not None:\n            # Account for a placeholder in the new graph.\n            node.args = (new_input_node,)\n            new_input_node = None\n        if node.op in [\"placeholder\", \"get_attr\", \"call_function\", \"call_method\", \"call_module\"]:\n            # Copy the nodes from the existing graph to the new graph.\n            new_node = new_graph.node_copy(node, lambda x: env[x.name])\n            env[node.name] = new_node\n        elif node.op == \"output\":\n            # If this is the last node, we should add an output\n            # node and add the last graph to the list.\n            assert prev_node, \"prev_node cannot be None\"\n\n            with new_graph.inserting_after(prev_node):\n                new_graph.output(env[prev_node.name])\n            module_list.append(torch.fx.GraphModule(model, new_graph))\n            break\n        prev_node = new_node\n        prev_shard_id = node_name_to_shard_id[node.name]\n\n    return module_list\n", "path": "fairscale/experimental/nn/auto_shard.py"}]}