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

Code before:
        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:
                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)

                param = local_state[key]

                try:
                    with torch.no_grad():
                        param.copy_(input_param)
                except Exception as ex:
                    error_msgs.append(
                        'While copying the parameter named "{}", '
                        "whose dimensions in the model are {} and "

Review comment (suggestion):
You could remove key assignment and directly use it later in `local_state.get`, though it's still sort of hard-coding...

Code after:
    ):
        # 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)

                try:
                    with torch.no_grad():
                        param.copy_(input_param)
                except Exception as ex:
                    error_msgs.append(
                        'While copying the parameter named "{}", '
                        "whose dimensions in the model are {} and "
                        "whose dimensions in the checkpoint are {}, "
                        "an exception occurred : {}.".format(key, param.size(), input_param.size(), ex.args)
                    )

                del self.gate_up_dict

            strict = False  # to avoid unexpected_keys
        super()._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        Args: