# Buggy code
def DistributedFairseqModel(args, model, process_group):
    """
    Wrap a *model* to support distributed data parallel training.

    This is similar to the built-in DistributedDataParallel, but allows
    additional configuration of the DistributedDataParallel class to
    use, and also provides easier access to the wrapped model by
    forwarding requests for missing attributes to the wrapped model.

    Args:
        args (argparse.Namespace): fairseq args
        model (BaseFairseqModel): model to wrap
        process_group: the c10d process group to be used for distributed data
            parallel all-reduction.
    """
    # determine which DDP class to extend
    assert isinstance(model, nn.Module)
    if args.tpu:
        ddp_class = TPUDistributedDataParallel
        init_kwargs = dict(
            module=model,
            process_group=process_group,
        )
    elif args.distributed_wrapper == "DDP" and args.ddp_backend == "c10d":
        ddp_class = nn.parallel.DistributedDataParallel
        init_kwargs = dict(
            module=model,
            device_ids=[args.device_id],
            output_device=args.device_id,
            broadcast_buffers=args.broadcast_buffers,
            bucket_cap_mb=args.bucket_cap_mb,
            process_group=process_group,
        )
        # Maintain backward compatibility
        if "check_reduction" in inspect.getargspec(ddp_class)[0]:
            init_kwargs["check_reduction"] = True
        if "find_unused_parameters" in inspect.getargspec(ddp_class)[0]:
            init_kwargs["find_unused_parameters"] = args.find_unused_parameters
    elif args.distributed_wrapper == "DDP" and args.ddp_backend == "no_c10d":
        ddp_class = LegacyDistributedDataParallel
        init_kwargs = dict(
            module=model,
            buffer_size=2**28,
            process_group=process_group,
        )
    elif args.distributed_wrapper == "SlowMo":
        if _GOSSIP_DISABLED:
            raise ImportError(
                "Cannot find gossip library. Please install from: "
                "github.com/facebookresearch/stochastic_gradient_push"
            )
        ddp_class = gossip.GossipDataParallel

        # The values of slowmo_momentum below were obtained by tuning on the
        # En-De 16 dataset by training the transformer_wmt_en_de_large model
        if args.slowmo_momentum is None:
            if args.distributed_world_size <= 16:
                args.slowmo_momentum = 0.0
            elif args.distributed_world_size <= 32:
                args.slowmo_momentum = 0.2
            elif args.distributed_world_size <= 64:
                args.slowmo_momentum = 0.5
            else:
                args.slowmo_momentum = 0.6

        init_kwargs = dict(
            module=model,
            device_ids=[args.device_id],
            output_device=args.device_id,
            broadcast_buffers=args.broadcast_buffers,
            nprocs_per_node=args.nprocs_per_node,
            slowmo_momentum=args.slowmo_momentum,
            localsgd=(args.slowmo_algorithm == "LocalSGD"),
            localsgd_frequency=args.localsgd_frequency,
        )
    else:
        raise ValueError("Unknown --ddp-backend: " + args.ddp_backend)

    class _DistributedFairseqModel(ddp_class):
        """Extend DistributedDataParallel to check for missing
        attributes in the wrapped module."""

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)

        def __getattr__(self, name):
            wrapped_module = super().__getattr__("module")
            if hasattr(wrapped_module, name):
                return getattr(wrapped_module, name)
            return super().__getattr__(name)

    return _DistributedFairseqModel(**init_kwargs)


# Traceback (AttributeError)
Traceback (most recent call last):
File "xxx/fairseq/fairseq/modules/multihead_attention.py", line 56, in __init__
self.reset_parameters()
File "xxx/fairseq/fairseq/modules/multihead_attention.py", line 82, in reset_parameters
nn.init.constant_(self.out_proj.bias, 0.)
File "xxx/torch/nn/init.py", line 120, in constant_
return _no_grad_fill_(tensor, val)
File "xxx/torch/nn/init.py", line 24, in _no_grad_fill_
return tensor.fill_(val)
AttributeError: 'NoneType' object has no attribute 'fill_'

# Fixed code
def DistributedFairseqModel(args, model, process_group):
    """
    Wrap a *model* to support distributed data parallel training.

    This is similar to the built-in DistributedDataParallel, but allows
    additional configuration of the DistributedDataParallel class to
    use, and also provides easier access to the wrapped model by
    forwarding requests for missing attributes to the wrapped model.

    Args:
        args (argparse.Namespace): fairseq args
        model (BaseFairseqModel): model to wrap
        process_group: the c10d process group to be used for distributed data
            parallel all-reduction.
    """
    # determine which DDP class to extend
    assert isinstance(model, nn.Module)
    if args.tpu:
        ddp_class = TPUDistributedDataParallel
        init_kwargs = dict(
            module=model,
            process_group=process_group,
        )
    elif args.distributed_wrapper == "DDP" and args.ddp_backend == "c10d":
        ddp_class = nn.parallel.DistributedDataParallel
        init_kwargs = dict(
            module=model,
            device_ids=[args.device_id],
            output_device=args.device_id,
            broadcast_buffers=args.broadcast_buffers,
            bucket_cap_mb=args.bucket_cap_mb,
            process_group=process_group,
        )
        # Maintain backward compatibility
        if "check_reduction" in inspect.getargspec(ddp_class)[0]:
            init_kwargs["check_reduction"] = True
        if "find_unused_parameters" in inspect.getargspec(ddp_class)[0]:
            init_kwargs["find_unused_parameters"] = args.find_unused_parameters
    elif args.distributed_wrapper == "DDP" and args.ddp_backend == "no_c10d":
        ddp_class = LegacyDistributedDataParallel
        init_kwargs = dict(
            module=model,
            buffer_size=2**28,
            process_group=process_group,
        )
    elif args.distributed_wrapper == "SlowMo":
        if _GOSSIP_DISABLED:
            raise ImportError(
                "Cannot find gossip library. Please install from: "
                "github.com/facebookresearch/stochastic_gradient_push"
            )
        ddp_class = gossip.GossipDataParallel

        # The values of slowmo_momentum below were obtained by tuning on the
        # En-De 16 dataset by training the transformer_wmt_en_de_large model
        if args.slowmo_momentum is None:
            if args.distributed_world_size <= 16:
                args.slowmo_momentum = 0.0
            elif args.distributed_world_size <= 32:
                args.slowmo_momentum = 0.2
            elif args.distributed_world_size <= 64:
                args.slowmo_momentum = 0.5
            else:
                args.slowmo_momentum = 0.6

        init_kwargs = dict(
            module=model,
            device_ids=[args.device_id],
            output_device=args.device_id,
            broadcast_buffers=args.broadcast_buffers,
            nprocs_per_node=args.nprocs_per_node,
            slowmo_momentum=args.slowmo_momentum,
            localsgd=(args.slowmo_algorithm == "LocalSGD"),
            localsgd_frequency=args.localsgd_frequency,
        )
    else:
        raise ValueError("Unknown --ddp-backend: " + args.ddp_backend)

    heartbeat_timeout = getattr(args, "heartbeat_timeout", -1)

    class _DistributedFairseqModel(ddp_class):
        """
        Extend DistributedDataParallel to check for missing attributes in the
        wrapped module and to add a timeout to kill the job if no progress is
        made (--heartbeat-timeout).
        """

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self._heartbeat_timeout = heartbeat_timeout
            if self._heartbeat_timeout > 0:
                self._heartbeat = threading.Event()
                self._heartbeat_thread = threading.Thread(
                    target=self._check_heartbeat,
                    args=(os.getpid(),),
                    daemon=True,
                )
                self._heartbeat_thread.start()
            else:
                self._heartbeat = None

        def _check_heartbeat(self, parent_pid):
            self._heartbeat.wait()  # wait for the first forward pass
            while True:
                self._heartbeat.clear()
                success = self._heartbeat.wait(timeout=self._heartbeat_timeout)
                if not success:
                    logger.error(
                        (
                            "Killing job for not making progress in {} seconds. "
                            "Set --heartbeat-timeout=-1 to disable this timeout."
                        ).format(int(self._heartbeat_timeout))
                    )
                    os.kill(parent_pid, signal.SIGKILL)
                    return

        def __getattr__(self, name):
            wrapped_module = super().__getattr__("module")
            if hasattr(wrapped_module, name):
                return getattr(wrapped_module, name)
            return super().__getattr__(name)

        def forward(self, *args, **kwargs):
            if self._heartbeat is not None:
                self._heartbeat.set()
            return super().forward(*args, **kwargs)

    return _DistributedFairseqModel(**init_kwargs)