We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
_amp_state determines whether running in distributed at import
It looks like the `_amp_state.py` module determines whether pytorch is running in distributed mode at the import level. The `distributed` only seems to be used in `maybe_print`. See code snippet:

https://github.com/NVIDIA/apex/blob/37cdaf4ad57ab4e7dd9ef13dbed7b29aa939d061/apex/amp/_amp_state.py#L38-L52

This causes a couple issues:

1. It will only support the `env://` initialization of torch distributed
2. It will fail if amp is imported before launching the distributed training

Neither of these is an issue for most, since most people launch via `torch.distributed.launch`. However, it can be an issue if you define your own distributed launch function or use `torch.multiprocessing.spawn`. I can't see a good reason to do it this way anyway, as it appears this variable is only used in the `maybe_print` function. I'll submit a pull request to fix this. Let me know if I'm missing something though.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `apex/amp/_amp_state.py`
Content:
```
1 # This is a "header object" that allows different amp modules to communicate.
2 # I'm a C++ guy, not a python guy.  I decided this approach because it seemed most C++-like.  
3 # But apparently it's ok:
4 # http://effbot.org/pyfaq/how-do-i-share-global-variables-across-modules.htm
5 import os
6 import torch
7 
8 TORCH_MAJOR = int(torch.__version__.split('.')[0])
9 TORCH_MINOR = int(torch.__version__.split('.')[1])
10 
11 if TORCH_MAJOR == 0:
12     import collections.abc as container_abcs
13 else:
14     from torch._six import container_abcs
15 
16 
17 class AmpState(object):
18     def __init__(self):
19         self.hard_override=False
20         self.allow_incoming_model_not_fp32 = False
21         self.verbosity=1
22 
23 
24 # Attribute stash.  Could also just stash things as global module attributes.
25 _amp_state = AmpState()
26 
27 
28 def warn_or_err(msg):
29     if _amp_state.hard_override:
30         print("Warning:  " + msg)
31     else:
32         raise RuntimeError(msg)
33         # I'm not sure if allowing hard_override is a good idea.
34         # + "  If you're sure you know what you're doing, supply " +
35         #                    "hard_override=True to amp.initialize.")
36 
37 
38 distributed = False
39 if 'WORLD_SIZE' in os.environ:
40     distributed = int(os.environ['WORLD_SIZE']) > 1
41 
42 
43 def maybe_print(msg, rank0=False):
44     if _amp_state.verbosity > 0:
45         if rank0:
46             if distributed:
47                 if torch.distributed.get_rank() == 0:
48                     print(msg)
49             else:
50                 print(msg)
51         else:
52             print(msg)
53 
54 
55 # def iter_params(param_groups):
56 #     for group in param_groups:
57 #         for p in group['params']:
58 #             yield p
59 
60 
61 def master_params(optimizer):
62     """
63     Generator expression that iterates over the params owned by ``optimizer``.
64 
65     Args:
66         optimizer: An optimizer previously returned from ``amp.initialize``.
67     """
68     for group in optimizer.param_groups:
69         for p in group['params']:
70             yield p
71
```
--- 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/apex/amp/_amp_state.py b/apex/amp/_amp_state.py
--- a/apex/amp/_amp_state.py
+++ b/apex/amp/_amp_state.py
@@ -1,5 +1,5 @@
 # This is a "header object" that allows different amp modules to communicate.
-# I'm a C++ guy, not a python guy.  I decided this approach because it seemed most C++-like.  
+# I'm a C++ guy, not a python guy.  I decided this approach because it seemed most C++-like.
 # But apparently it's ok:
 # http://effbot.org/pyfaq/how-do-i-share-global-variables-across-modules.htm
 import os
@@ -35,12 +35,9 @@
         #                    "hard_override=True to amp.initialize.")
 
 
-distributed = False
-if 'WORLD_SIZE' in os.environ:
-    distributed = int(os.environ['WORLD_SIZE']) > 1
-
-
 def maybe_print(msg, rank0=False):
+    distributed = torch.distributed.is_initialized() and \
+        torch.distributed.get_world_size() > 1
     if _amp_state.verbosity > 0:
         if rank0:
             if distributed:

{"golden_diff": "diff --git a/apex/amp/_amp_state.py b/apex/amp/_amp_state.py\n--- a/apex/amp/_amp_state.py\n+++ b/apex/amp/_amp_state.py\n@@ -1,5 +1,5 @@\n # This is a \"header object\" that allows different amp modules to communicate.\n-# I'm a C++ guy, not a python guy.  I decided this approach because it seemed most C++-like.  \n+# I'm a C++ guy, not a python guy.  I decided this approach because it seemed most C++-like.\n # But apparently it's ok:\n # http://effbot.org/pyfaq/how-do-i-share-global-variables-across-modules.htm\n import os\n@@ -35,12 +35,9 @@\n         #                    \"hard_override=True to amp.initialize.\")\n \n \n-distributed = False\n-if 'WORLD_SIZE' in os.environ:\n-    distributed = int(os.environ['WORLD_SIZE']) > 1\n-\n-\n def maybe_print(msg, rank0=False):\n+    distributed = torch.distributed.is_initialized() and \\\n+        torch.distributed.get_world_size() > 1\n     if _amp_state.verbosity > 0:\n         if rank0:\n             if distributed:\n", "issue": "_amp_state determines whether running in distributed at import\nIt looks like the `_amp_state.py` module determines whether pytorch is running in distributed mode at the import level. The `distributed` only seems to be used in `maybe_print`. See code snippet:\r\n\r\nhttps://github.com/NVIDIA/apex/blob/37cdaf4ad57ab4e7dd9ef13dbed7b29aa939d061/apex/amp/_amp_state.py#L38-L52\r\n\r\nThis causes a couple issues:\r\n\r\n1. It will only support the `env://` initialization of torch distributed\r\n2. It will fail if amp is imported before launching the distributed training\r\n\r\nNeither of these is an issue for most, since most people launch via `torch.distributed.launch`. However, it can be an issue if you define your own distributed launch function or use `torch.multiprocessing.spawn`. I can't see a good reason to do it this way anyway, as it appears this variable is only used in the `maybe_print` function. I'll submit a pull request to fix this. Let me know if I'm missing something though.\n", "before_files": [{"content": "# This is a \"header object\" that allows different amp modules to communicate.\n# I'm a C++ guy, not a python guy.  I decided this approach because it seemed most C++-like.  \n# But apparently it's ok:\n# http://effbot.org/pyfaq/how-do-i-share-global-variables-across-modules.htm\nimport os\nimport torch\n\nTORCH_MAJOR = int(torch.__version__.split('.')[0])\nTORCH_MINOR = int(torch.__version__.split('.')[1])\n\nif TORCH_MAJOR == 0:\n    import collections.abc as container_abcs\nelse:\n    from torch._six import container_abcs\n\n\nclass AmpState(object):\n    def __init__(self):\n        self.hard_override=False\n        self.allow_incoming_model_not_fp32 = False\n        self.verbosity=1\n\n\n# Attribute stash.  Could also just stash things as global module attributes.\n_amp_state = AmpState()\n\n\ndef warn_or_err(msg):\n    if _amp_state.hard_override:\n        print(\"Warning:  \" + msg)\n    else:\n        raise RuntimeError(msg)\n        # I'm not sure if allowing hard_override is a good idea.\n        # + \"  If you're sure you know what you're doing, supply \" +\n        #                    \"hard_override=True to amp.initialize.\")\n\n\ndistributed = False\nif 'WORLD_SIZE' in os.environ:\n    distributed = int(os.environ['WORLD_SIZE']) > 1\n\n\ndef maybe_print(msg, rank0=False):\n    if _amp_state.verbosity > 0:\n        if rank0:\n            if distributed:\n                if torch.distributed.get_rank() == 0:\n                    print(msg)\n            else:\n                print(msg)\n        else:\n            print(msg)\n\n\n# def iter_params(param_groups):\n#     for group in param_groups:\n#         for p in group['params']:\n#             yield p\n\n\ndef master_params(optimizer):\n    \"\"\"\n    Generator expression that iterates over the params owned by ``optimizer``.\n\n    Args:\n        optimizer: An optimizer previously returned from ``amp.initialize``.\n    \"\"\"\n    for group in optimizer.param_groups:\n        for p in group['params']:\n            yield p\n", "path": "apex/amp/_amp_state.py"}], "after_files": [{"content": "# This is a \"header object\" that allows different amp modules to communicate.\n# I'm a C++ guy, not a python guy.  I decided this approach because it seemed most C++-like.\n# But apparently it's ok:\n# http://effbot.org/pyfaq/how-do-i-share-global-variables-across-modules.htm\nimport os\nimport torch\n\nTORCH_MAJOR = int(torch.__version__.split('.')[0])\nTORCH_MINOR = int(torch.__version__.split('.')[1])\n\nif TORCH_MAJOR == 0:\n    import collections.abc as container_abcs\nelse:\n    from torch._six import container_abcs\n\n\nclass AmpState(object):\n    def __init__(self):\n        self.hard_override=False\n        self.allow_incoming_model_not_fp32 = False\n        self.verbosity=1\n\n\n# Attribute stash.  Could also just stash things as global module attributes.\n_amp_state = AmpState()\n\n\ndef warn_or_err(msg):\n    if _amp_state.hard_override:\n        print(\"Warning:  \" + msg)\n    else:\n        raise RuntimeError(msg)\n        # I'm not sure if allowing hard_override is a good idea.\n        # + \"  If you're sure you know what you're doing, supply \" +\n        #                    \"hard_override=True to amp.initialize.\")\n\n\ndef maybe_print(msg, rank0=False):\n    distributed = torch.distributed.is_initialized() and \\\n        torch.distributed.get_world_size() > 1\n    if _amp_state.verbosity > 0:\n        if rank0:\n            if distributed:\n                if torch.distributed.get_rank() == 0:\n                    print(msg)\n            else:\n                print(msg)\n        else:\n            print(msg)\n\n\n# def iter_params(param_groups):\n#     for group in param_groups:\n#         for p in group['params']:\n#             yield p\n\n\ndef master_params(optimizer):\n    \"\"\"\n    Generator expression that iterates over the params owned by ``optimizer``.\n\n    Args:\n        optimizer: An optimizer previously returned from ``amp.initialize``.\n    \"\"\"\n    for group in optimizer.param_groups:\n        for p in group['params']:\n            yield p\n", "path": "apex/amp/_amp_state.py"}]}