AttributeError: Can't get attribute 'MegatronPolicy' on <module '__main__' (built-in)>

[Oaklight]

Trying to use parallelformers with the megatron-11b pip package. The MegatronPolicy class is as-provided from megatron-11b pypi webpage

How to reproduce

from megatron_11b import MegatronForCausalLM, MegatronTokenizer

tokenizer = MegatronTokenizer.from_pretrained("./megatron-11B")
model = MegatronForCausalLM.from_pretrained("./megatron-11B")

# https://tunib-ai.github.io/parallelformers/intro/POLICY.html

from parallelformers.policies.base import Policy, Layer
from parallelformers.utils.dist_utils import AllReduceLinear
from megatron_11b.modeling_megatron import MegatronDecoderLayer

class MegatronPolicy(Policy):

    @staticmethod
    def replace_arguments(config, world_size):
        return {
            # 1. reduce hidden size
            "self_attn.embed_dim": config.d_model // world_size,

            # 2. reduce number of heads
            "self_attn.num_heads": config.encoder_attention_heads // world_size,
        }

    @staticmethod
    def attn_qkv():
        return [
            Layer(
                weight="self_attn.q_proj.weight",
                bias="self_attn.q_proj.bias",
            ),
            Layer(
                weight="self_attn.k_proj.weight",
                bias="self_attn.k_proj.bias",
            ),
            Layer(
                weight="self_attn.v_proj.weight",
                bias="self_attn.v_proj.bias",
            ),
        ]

    @staticmethod
    def attn_out():
        return [
            Layer(
                weight="self_attn.out_proj.weight",
                bias="self_attn.out_proj.bias",
                replace=AllReduceLinear,
            ),
        ]

    @staticmethod
    def mlp_in():
        return [
            Layer(
                weight="fc1.weight",
                bias="fc1.bias",
            ),
        ]

    @staticmethod
    def mlp_out():
        return [
            Layer(
                weight="fc2.weight",
                bias="fc2.bias",
                replace=AllReduceLinear,
            ),
        ]

    @staticmethod
    def original_layer_class():
        return MegatronDecoderLayer

from parallelformers import parallelize

parallelize(model, num_gpus=8, fp16=True, verbose='detail', custom_policies=[MegatronPolicy])

Environment

• OS : Ubuntu LTS 20.04
• Python version : 3.8
• Transformers version : 4.4.2
• Whether to use Docker: no
• Misc.: it's executed in a jupyter notebook, which might be the source of the problem: https://stackoverflow.com/a/65001152

[Oaklight]

after some digging, it's very likely caused by the jupyter/ipython's handling of spawn: Similar situations got discussed in several places. https://github.com/pytorch/pytorch/issues/17680#issuecomment-491384202 for example. I'm testing the above code in .py script, and will try to modify parallelformer's parallelize.py to replace torch.multiprocessing with multiprocess

[Oaklight]

unfortunately, directly running the above code in script mode throw this error

$ python megatron-11b-prompt.py

Traceback (most recent call last):
  File "/homes/pding/miniconda3/envs/megatron/lib/python3.8/site-packages/parallelformers/parallelize.py", line 283, in parallelize
    process.start()
  File "/homes/pding/miniconda3/envs/megatron/lib/python3.8/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/homes/pding/miniconda3/envs/megatron/lib/python3.8/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/homes/pding/miniconda3/envs/megatron/lib/python3.8/multiprocessing/context.py", line 283, in _Popen
    return Popen(process_obj)
  File "/homes/pding/miniconda3/envs/megatron/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/homes/pding/miniconda3/envs/megatron/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/homes/pding/miniconda3/envs/megatron/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 42, in _launch
    prep_data = spawn.get_preparation_data(process_obj._name)
  File "/homes/pding/miniconda3/envs/megatron/lib/python3.8/multiprocessing/spawn.py", line 154, in get_preparation_data
    _check_not_importing_main()
  File "/homes/pding/miniconda3/envs/megatron/lib/python3.8/multiprocessing/spawn.py", line 134, in _check_not_importing_main
    raise RuntimeError('''
RuntimeError:
        An attempt has been made to start a new process before the
        current process has finished its bootstrapping phase.

        This probably means that you are not using fork to start your
        child processes and you have forgotten to use the proper idiom
        in the main module:

            if __name__ == '__main__':
                freeze_support()
                ...

        The "freeze_support()" line can be omitted if the program
        is not going to be frozen to produce an executable.

[Oaklight]

ok, fixed by replacing torch.multiprocessing with multiprocess in parallelize.py PR open at #21

[hyunwoongko]

Didn't these changes cause any other problems?

[Oaklight]

I didn't see other error

[hyunwoongko]

Ok I see. thank you.