prompt: 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()
  
  ```
  
golden_diff: 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:
  
verification_info: {"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"}]}