Support OLMo models.

[Isotr0py]

Related issue:

• Support for new OLMo models #2763

TODO:

• Add model config OLMoConfig
• Test on OLMo-1B
• Test on OLMo-7B/7B-Twin-2T
• Format code

This is still in progress before all developments and tests finish.
Done.

[Isotr0py]

It seems that SwiGLU activation used in olmo is different from the SiluAndMul in vllm:

class SiluAndMul(nn.Module):
    """An activation function for SwiGLU.

    The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.

    Shapes:
        x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d)
        return: (batch_size, seq_len, d) or (num_tokens, d)
    """

    def _forward(self, x: torch.Tensor) -> torch.Tensor:
        """PyTorch-native implementation equivalent to forward()."""
        d = x.shape[-1] // 2
        return F.silu(x[..., :d]) * x[..., d:]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        d = x.shape[-1] // 2
        output_shape = (x.shape[:-1] + (d, ))
        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
        ops.silu_and_mul(out, x)
        return out

[zhuohan123]

LGTM! Thank you for your contribution!