Repository: hpcaitech/ColossalAI
File: colossalai/inference/modeling/models/nopadding_llama.py (Python)

Code before:
            "up_proj.weight": None,
        }  # used and delattr in load/shard of gate/up weight
        self.down_proj = mlp_dproj
        self.process_group = process_group

    @staticmethod
    def from_native_module(
        module: LlamaMLP, process_group: Union[ProcessGroup, List[ProcessGroup]], *args, **kwargs
    ) -> ParallelModule:
        """Used for initialize the weight of NopadLlamaMLP by origin LlamaMLP.

        Args:
            module (LlamaMLP): The origin LlamaMLP layer.
        """

        config = module.config

        mlp_gproj_w = module.gate_proj.weight
        assert is_distributed_tensor(
            module.gate_proj.weight
        ), "gate_proj.weight must be dtensor so we could get the layout of the weight"
        mlp_uproj_w = module.up_proj.weight
        mlp_dproj = module.down_proj

        mlp_layer = NopadLlamaMLP(
            config=config,
            mlp_gproj_w=mlp_gproj_w,
            mlp_uproj_w=mlp_uproj_w,
            mlp_dproj=mlp_dproj,
            process_group=process_group,
        )

        return mlp_layer

    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
        # NOTE This is a hack to ensure we could load the right weight from LlamaMLP checkpoint due to the use of torch.stack(gate_weight, up_weight)

        if hasattr(self, "gate_up_dict"):
            for hook in self._load_state_dict_pre_hooks.values():
                hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)

            persistent_buffers = {k: v for k, v in self._buffers.items() if k not in self._non_persistent_buffers_set}
            local_name_params = itertools.chain(self._parameters.items(), persistent_buffers.items())
            local_state = {k: v for k, v in local_name_params if v is not None}

            key = "gate_up_weight"
            device_mesh = self.helper_layout.device_mesh
            sharding_spec = self.helper_layout.sharding_spec
            for weight_name in self.gate_up_dict:

Review comment (suggestion):
1. With `self.gate_up_dict` already set in `__init__`, the if block is noneffective and redundant. 
2. prevent hard-coding

Code after:

        config = module.config

        mlp_gproj_w = module.gate_proj.weight
        assert is_distributed_tensor(
            module.gate_proj.weight
        ), "gate_proj.weight must be dtensor so we could get the layout of the weight"
        mlp_uproj_w = module.up_proj.weight
        mlp_dproj = module.down_proj

        mlp_layer = NopadLlamaMLP(
            config=config,
            mlp_gproj_w=mlp_gproj_w,
            mlp_uproj_w=mlp_uproj_w,
            mlp_dproj=mlp_dproj,
            process_group=process_group,
        )

        return mlp_layer

    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
        # NOTE This is a hack to ensure we could load the right weight from LlamaMLP checkpoint due to the use of torch.stack(gate_weight, up_weight)

        if hasattr(self, "gate_up_dict"):
            for hook in self._load_state_dict_pre_hooks.values():
                hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)

            persistent_buffers = {k: v for k, v in self._buffers.items() if k not in self._non_persistent_buffers_set}
            local_name_params = itertools.chain(self._parameters.items(), persistent_buffers.items())
            local_state = {k: v for k, v in local_name_params if v is not None}

            device_mesh = self.helper_layout.device_mesh
            sharding_spec = self.helper_layout.sharding_spec
            for weight_name in self.gate_up_dict:
                prefix_weight_name = prefix + weight_name
                if prefix_weight_name in state_dict.keys():
                    w = distribute_tensor(state_dict[prefix_weight_name], device_mesh, sharding_spec)
                    self.gate_up_dict[weight_name] = w.T

            if None not in self.gate_up_dict.values():
                # we've got all the weights of gate/up
                gate_up_w = torch.stack(list(self.gate_up_dict.values()), dim=0)

                input_param = nn.Parameter(
                    gate_up_w
                )  # NOTE gate_up_weight doesn't have to be a distensor, Like input_param = sharded_tensor_to_param(input_param)

                key = "gate_up_weight"
                param = local_state.get(key, None)