Add Phi-4 Backbone #2272

yrahul3910 · 2025-05-27T04:35:06Z

Description of the change

This is the first PR in contributing the Phi-4 model to KerasHub, and includes the backbone and its test file.

Reference

Colab Notebook

I've had some trouble getting this part to work, so I need some help. This is my Colab notebook, but the HF model has been pretty annoying to run. On CPU machines, it seems to constantly allocate all available memory (I gave up after giving it 280GB), and on an H200 on Modal, I couldn't get an output after 15 minutes. In the notebook, this line:

hf_output = pt_model(**hf_sample_input)

at the bottom is the one I have trouble with.

Checklist

I have added all the necessary unit tests for my change.
I have verified that my change does not break existing code and works with all backends (TensorFlow, JAX, and PyTorch).
My PR is based on the latest changes of the main branch (if unsure, rebase the code).
I have followed the Keras Hub Model contribution guidelines in making these changes.
I have followed the Keras Hub API design guidelines in making these changes.
I have signed the Contributor License Agreement.

yrahul3910 · 2025-05-27T04:39:06Z

I uploaded the output from the KerasHub model, could someone upload a HF version that I can compare and add to the Colab?

yrahul3910 · 2025-05-27T04:41:44Z

P.S. A lot of this code is based on the existing code for Phi-3 (the technical report states it mostly follows the Phi-3-Medium architecture; I simply made the changes from the report and the reference implementation). Should I refactor it to inherit from Phi3Backbone?

sachinprasadhs

Thanks for the PR! added my review comments.

sachinprasadhs · 2025-05-30T18:00:35Z

keras_hub/src/models/phi4/phi4_attention.py

+                '`rope_scaling_type` must be `None` or `"su"`.'
+                "if `None` is choosed, `RotaryEmbedding` will be used."
+                'if `"su"` is choosed, `Phi4SuScaledRotaryEmbedding` will be '
+                "used."


May be change this to --> "rope_scaling_type must be None or su. If None, RotaryEmbedding will be used. If su, Phi4SuScaledRotaryEmbedding will be used."
Add backtick wherever it is necessary.

sachinprasadhs · 2025-05-30T18:48:47Z

keras_hub/src/models/phi4/phi4_backbone.py

+        vocabulary_size (int): The size of the token vocabulary. Defaults to
+            `100_352`.


Change this to --> vocabulary_size: int. The size of the token vocabulary. Defaults to `100_352`.

Follow the above arg pattern for others as well, i know this follows same as phi3, but this will be consistent with majority of our models.

sachinprasadhs · 2025-05-30T20:01:52Z

keras_hub/src/models/phi4/phi4_backbone_test.py

+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        for preset in Phi4Backbone.presets:
+            self.run_preset_test(
+                cls=Phi4Backbone,
+                preset=preset,
+                input_data=self.input_data,
+            )


Usually how big these models will be and how many presets are we testing here?

sachinprasadhs · 2025-05-30T20:03:54Z

keras_hub/src/models/phi4/phi4_decoder.py

+        hidden_dim=5120,
+        intermediate_dim=17_920,
+        num_query_heads=40,
+        num_key_value_heads=10,
+        activation="silu",
+        layer_norm_epsilon=1e-5,
+        kernel_initializer="glorot_uniform",
+        dropout=0,
+        max_sequence_length=16_384,
+        pretraining_sequence_length=16_384,
+        rope_max_wavelength=250_000,


Are these default values are mostly common for all the presets in phi-4, if not may be we can remove default values?

sachinprasadhs · 2025-05-30T20:06:46Z

keras_hub/src/models/phi4/phi4_layernorm.py

+# TODO: Deprecate this in favor of
+# `keras.layers.LayerNormalization(rms_scaling=True)` once Keras 2 support is
+# removed.


We don't have keras 2 support now, either update the code or remove/update this comment.

sachinprasadhs · 2025-05-30T20:11:56Z

P.S. A lot of this code is based on the existing code for Phi-3 (the technical report states it mostly follows the Phi-3-Medium architecture; I simply made the changes from the report and the reference implementation). Should I refactor it to inherit from Phi3Backbone?

How much similar is Phi-4 compared to Phi-3? What is the approx percentage of code we can reuse?

sachinprasadhs · 2025-05-30T20:55:30Z

I guess still Tokenizer and CausalLM and preset file with necessary test files still needs to be added?

yrahul3910 · 2025-06-01T19:38:56Z

Actually, I think we might get away with directly subclassing Phi3Backbone and changing the defaults. If you do

diff src/models/phi3/phi3_backbone.py src/models/phi4/phi4_backbone.py

The only differences are the model name and the defaults; initially I did this copy anticipating architectural changes, but it seems the only ones are in the attention. From the paper's Section 3:

The architecture closely follows phi-3-medium, except that we now use the tiktoken tokenizer (for better
multilingual support) with a padded vocabulary size of 100,352 (including unused tokens) and we use
full attention over the 4K context length, rather than a 2K sliding window used in phi-3-medium.

I could not find this sliding window attention in the code, however, so that also remained unchanged, and the tokenizer would be part of the third PR (based on the contributing guidelines). Do you think it's better if I just did that instead?

mattdangerw · 2025-06-13T19:35:51Z

Yeah this is an interesting question if the only thing that is really changing is the tokenizer. Would it work to just subclass all the phi3 classes as stubs that only update the classes? Like this

@keras_hub_export("keras_hub.models.Phi4CausalLM")
class Phi4CausalLM(Phi3CausalLM):
    backbone_cls = Phi4Backbone
    preprocessor_cls = Phi4CausalLMPreprocessor

And then define a new tokenizer? I don't think we need to worry about switchin the defaults, that we can just reflect in the preset configs we upload to kaggle.

Might be worth trying all of that on a single PR and trying to convert a model to see if everything works. Since we'd mostly be dealing with stub classes it shouldn't be too much code.

yrahul3910 · 2025-06-13T20:36:31Z

Yes, that makes sense to me. I'll try getting it out over the weekend.

abheesht17 · 2025-07-10T19:07:01Z

@yrahul3910 - are you still working on this?

yrahul3910 · 2025-07-10T20:12:49Z

Yes, so sorry for the delay! I'll push some commits soon.

divyashreepathihalli · 2025-07-10T23:59:53Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces the Phi-4 backbone and its associated components. The implementation is solid and well-structured. I've identified a few minor areas for improvement, mainly related to typos and redundant code. Addressing these will improve the clarity and maintainability of the new model.

keras_hub/src/models/phi4/phi4_backbone.py

keras_hub/src/models/phi4/phi4_attention.py

keras_hub/src/models/phi4/phi4_rotary_embedding.py

yrahul3910 · 2025-07-17T02:29:21Z

Sorry for the delays! I've taken @mattdangerw's advice and used subclasses instead. I'm a little less confident about the tokenizer bit--the paper mentions tiktoken, but since that's just a BPE tokenizer, I used the one we have in KerasHub and then followed the Llama-3 and BART model files (because Phi-4 ultimately is a modified version of Llama models, iirc, and these two use BPE as well).

ETA: I'm also not sure if I'll need to subclass and modify the attention layers, because I'm not quite sure where this part comes in (from the paper):

The architecture closely follows phi-3-medium, except that [...] we use full attention over the 4K context length, rather than a 2K sliding window used in phi-3-medium.

yrahul3910 · 2025-07-20T23:22:07Z

Fixed the tests! For some reason, the test_early_stopping test in Phi4CausalLMTest wouldn't work unless the pad token was at position 0, even if the pad_token_id was explicitly set correctly. This meant the padding mask was set to all True, when it should've had False at the last few positions (for example, Llama-3 does this correctly). So when we create the config files, we'll want to make sure this is true in the vocabulary

sachinprasadhs · 2025-08-13T04:27:50Z

@yrahul3910 , Let us know if this PR is ready for review.

yrahul3910 · 2025-08-13T23:40:40Z

Yes, it is!

yrahul3910 and others added 3 commits May 7, 2025 21:54

feat(phi4): add phi4_backbone

4a8566b

docs(phi4): update defaults in docstring

69f66ff

Merge branch 'keras-team:master' into master

1bfe756

divyashreepathihalli requested a review from sachinprasadhs May 28, 2025 02:47

sachinprasadhs reviewed May 30, 2025

View reviewed changes

sachinprasadhs added the stat:awaiting response from contributor label May 30, 2025

sachinprasadhs removed the stat:awaiting response from contributor label Jun 16, 2025

divyashreepathihalli added this to KerasHub Jul 10, 2025

divyashreepathihalli moved this to In Progress in KerasHub Jul 10, 2025

gemini-code-assist bot reviewed Jul 11, 2025

View reviewed changes

yrahul3910 added 3 commits July 16, 2025 22:27

feat(phi4): refactor Phi4Backbone to inherit from Phi-3

8b7146e

feat(phi4): add phi-4 tokenizer

3df73af

feat(phi4): add phi-4 causal_lm files

4aceea3

yrahul3910 added 4 commits July 16, 2025 22:34

fix(phi4): update docstring to use correct variable names

17a30ce

fix(phi4): remove dedicated attention and decoder modules

82b2912

fix(phi4): remove unused layernorm and rotary embedding layers

cbdf6ce

fix(phi4): fix unit tests

ce07951

fix(phi4): fix unit tests

0d13049

		vocabulary_size (int): The size of the token vocabulary. Defaults to
		`100_352`.

Add Phi-4 Backbone #2272

Are you sure you want to change the base?

Add Phi-4 Backbone #2272

Uh oh!

Conversation

yrahul3910 commented May 27, 2025

Description of the change

Reference

Colab Notebook

Checklist

Uh oh!

yrahul3910 commented May 27, 2025

Uh oh!

yrahul3910 commented May 27, 2025

Uh oh!

sachinprasadhs left a comment

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs May 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs May 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs May 30, 2025

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs May 30, 2025

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs May 30, 2025

Choose a reason for hiding this comment

Uh oh!

sachinprasadhs commented May 30, 2025

Uh oh!

sachinprasadhs commented May 30, 2025

Uh oh!

yrahul3910 commented Jun 1, 2025

Uh oh!

mattdangerw commented Jun 13, 2025

Uh oh!

yrahul3910 commented Jun 13, 2025

Uh oh!

abheesht17 commented Jul 10, 2025

Uh oh!

yrahul3910 commented Jul 10, 2025

Uh oh!

divyashreepathihalli commented Jul 10, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yrahul3910 commented Jul 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yrahul3910 commented Jul 20, 2025

Uh oh!

sachinprasadhs commented Aug 13, 2025

Uh oh!

yrahul3910 commented Aug 13, 2025

Uh oh!

Uh oh!

sachinprasadhs May 30, 2025 •

edited

Loading

sachinprasadhs May 30, 2025 •

edited

Loading

yrahul3910 commented Jul 17, 2025 •

edited

Loading