def get_past_kv_inputs(self, model_attributes, past_seq_len: int, use_fp16: bool, world_size: int = 1):
        """Get past_key_values for all layers.

        Shape of past_key_values is (num_heads, past_seq_len, head_size).
        """
        from olive.common.hf.wrapper import ModelWrapper

        model_wrapper = ModelWrapper(model_attributes)
        num_key_value_heads = model_wrapper.num_key_value_heads
        num_attention_heads = model_wrapper.num_attention_heads
        hidden_size = model_wrapper.hidden_size
        if num_attention_heads is None or hidden_size is None:
            raise ValueError("Cannot find num_attention_heads or hidden_size in model attributes")
        num_attention_heads = num_attention_heads // world_size

        head_size = hidden_size // num_attention_heads
        if num_key_value_heads is not None:
            num_key_value_heads = num_key_value_heads // world_size
            # adjust num_attention_heads to num_key_value_heads for MoE models to get the right shape
            num_attention_heads = num_key_value_heads

        num_hidden_layers = model_wrapper.num_hidden_layers
        torch_dtype = torch.float16 if use_fp16 else torch.float32
        return [
            (
                torch.rand(num_attention_heads, past_seq_len, head_size, dtype=torch_dtype),
                torch.rand(num_attention_heads, past_seq_len, head_size, dtype=torch_dtype),
            )
            for _ in range(num_hidden_layers)
        ]
Get past_key_values for all layers.

        Shape of past_key_values is (num_heads, past_seq_len, head_size).
        
olive/data/component/dataset.py