Staging PR for implimenting Phi-2 support.

[cm2435]

….org/main/getting-started/tutorials/05-layer-norm.html]

[cm2435]

Addresses #85

[danielhanchen]

@cm2435 Oh cool great work!! Was just working on the Jan release :) Ye using torch.allclose is what I normally use!!

[cm2435]

@danielhanchen I've implemented the ReLU and LayerNorm kernels for the model and added some test coverage around their closeness. Mind having a review when you get a min? ty!

[danielhanchen]

@cm2435 Great work again!! I'll do a PR review! :)

[danielhanchen]

@cm2435 Super great work! Some questions and notes from my side:

1. torch.manual_seed sadly only sets the seed on the CPU side - maybe you meant torch.cuda.manual_seed? I normally use transformers.set_seed
2. On ReLU - is this just an extra kernel you would like to have - or is this related to Phi? Phi I think uses gelu_new and not ReLU, and one issue I needed to solve was to find the derivative of Gelu. ReLU derivative is simple, since it's just either 1 or 0 for whether it was fired or not multiplied by the backprop deltas.
3. Triton's layernorm sadly is extremely cumbersome last I checked - https://github.com/lucidrains/triton-transformer/blob/main/triton_transformer/layernorm.py might be more useful - in fact Lucidrain's other kernels might be useful - I haven't checked them all yet though.
4. Another issue I need to investigate is what is "partial RoPE?" Phi has a scale of 40% 0.4:

# Partial rotary embedding
query_rot, query_pass = (
    query_states[..., : self.rotary_emb.dim],
    query_states[..., self.rotary_emb.dim :],
)
key_rot, key_pass = (
    key_states[..., : self.rotary_emb.dim],
    key_states[..., self.rotary_emb.dim :],
)
# [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor]
query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)

I'm assuming this is RoPE but only on the first 40% of head dimensions? Super weird - this'll affect the RoPE kernel as well.

All in all super good work - Phi just can be very annoying to work with, since i'm not sure why the authors of Phi decided to dramatically change a lot from Llama ie:

1. Gelu instead of Swiglu
2. Partial RoPE instead of full RoPE (does not making training that much slower? so unsure why they're doing this - Stability AI's Stable Code release also did this - unsure on their reasoning?)
3. Residual dropout after attention and MLP, ie dropout 10% of the attention output and MLP outputs - I guess this partially makes sense due to counteract overfitting.

[cm2435]

Yep. I'm dumb and misread GeLU for ReLU. I need more coffee in my life haha.

I'll take a crack this week at contributing a GeLU fwd and backward kernel, and tidy up the other bits.

Yeah, that partial rope embedding is weird, some GitHub posts I've found on it suggest it was done to apply the rotational transform to the first X % of the embedding dimension.

To quote the author in that issue:
"yup, we found that it worked very slightly better, never any worse, in the exception that the value is too low (I saw performance dip when frequencies reach around 6 - dimension of 24)"

It should affect any triton kernels for it, but it shouldn't be too hard, surely it's just an extra kernel parameter as to what fraction of the embedding to apply the transform to.

[cm2435]

As a separate point; what are you using as your reference implementation? Are you working off the paper or do you have code?
I've found transformers code but that's for phi-1 which seems to diverge a good chunk in architecture.

[danielhanchen]

@cm2435 Interesting on partial RoPE!! I did hear it was used in the new Stable Code quote https://huggingface.co/stabilityai/stable-code-3b

Position Embeddings: Rotary Position Embeddings (Su et al., 2021) applied to the first 25% of head embedding dimensions for improved throughput following Black et al. (2022).

Unsure on whether accuracy would improve - I actually thought this might make Phi-2 not RoPE scalable ie it cannot be adjusted "properly" to handle longer sequences sadly. Unsure though. Throughput - extremely unlikely an issue if one uses good kernels.

But I'll research more on increased accuracy - that's interesting! Implementation should be fine hopefully - can just take the view as what Phi-2 did.

In terms of reference implementation - transformers 4.37 got released - you can now follow their official implementation here: https://github.com/huggingface/transformers/blob/main/src/transformers/models/phi/modeling_phi.py

[danielhanchen]

@cm2435 So tried Wolfram but then found this paper: https://arxiv.org/pdf/2305.12073.pdf

Absolute horrible nightmare

[cm2435]

@danielhanchen Not had time to read that paper but I'd be careful because at a glance that looks like an approximation.

GeLU involves a $erf(x)$, those $tanh(...)$ terms are approximations
https://paperswithcode.com/method/gelu

I think it should be fine but that's not strictly the derivative. I've got a free hour or two so I'll tack a swing at an implementation

[danielhanchen]

@cm2435 Oh I don't think Phi uses the error function, but rather an approximation

[cm2435]

@danielhanchen That's true. I've taken a first punt at a forward and backward GeLU implementation but I still need to test it. After that will get onto either the partial rope factor or the dropout.

[danielhanchen]

@cm2435 Cool fabulous work!

[cm2435]

@danielhanchen It's Been a minute since I pushed something but some progress. I've patched the rope scaling kernel to take In a partial scaling factor (which I think is in line with what is done in phi? but you will absolutely want to check. I've also done a basic pass at a seeded dropout kernel based on Ludacrens implementation and the triton tutorials.

If those are good, based on the first message you sent I think we can move on to implementing the model?