Merge branch 'main' into tcc/bench_vis

This commit is contained in:
Shao Tang
2025-11-23 09:12:53 -06:00
committed by GitHub
28 changed files with 2877 additions and 124 deletions
+3
View File
@@ -262,8 +262,11 @@ loss.backward()
| Phi3 & Phi3.5 | `liger_kernel.transformers.apply_liger_kernel_to_phi3` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
| Granite 3.0 & 3.1 | `liger_kernel.transformers.apply_liger_kernel_to_granite` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss |
| OLMo2 | `liger_kernel.transformers.apply_liger_kernel_to_olmo2` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
| Olmo3 | `liger_kernel.transformers.apply_liger_kernel_to_olmo3` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
| GLM-4 | `liger_kernel.transformers.apply_liger_kernel_to_glm4` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
| InternVL3 | `liger_kernel.transformers.apply_liger_kernel_to_internvl` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
| HunyuanV1 | `liger_kernel.transformers.apply_liger_kernel_to_hunyuan_v1_dense` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
| HunyuanV1 MoE | `liger_kernel.transformers.apply_liger_kernel_to_hunyuan_v1_moe` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
## Low-level APIs
+240
View File
@@ -1703,3 +1703,243 @@ llama4_rope,huggingface,full,memory,MB,T,sequence length,2048,314.01611328125,31
llama4_rope,huggingface,full,memory,MB,T,sequence length,4096,596.03173828125,596.03173828125,596.03173828125,"{""dtype"": ""torch.bfloat16"", ""hidden_size"": 8192, ""num_q_heads"": 32, ""num_kv_heads"": 8}",NVIDIA H100 80GB HBM3,2025-08-07 21:42:21,0.6.1
llama4_rope,huggingface,full,memory,MB,T,sequence length,8192,1160.06298828125,1160.06298828125,1160.06298828125,"{""dtype"": ""torch.bfloat16"", ""hidden_size"": 8192, ""num_q_heads"": 32, ""num_kv_heads"": 8}",NVIDIA H100 80GB HBM3,2025-08-07 21:42:21,0.6.1
llama4_rope,huggingface,full,memory,MB,T,sequence length,16384,2288.12548828125,2288.12548828125,2288.12548828125,"{""dtype"": ""torch.bfloat16"", ""hidden_size"": 8192, ""num_q_heads"": 32, ""num_kv_heads"": 8}",NVIDIA H100 80GB HBM3,2025-08-07 21:42:21,0.6.1
tiled_geglu,liger,full,speed,ms,T,sequence length,1024,2.1678080558776855,2.166579246520996,2.1682305335998535,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:48,0.6.3
tiled_geglu,liger,full,speed,ms,T,sequence length,2048,4.344256401062012,4.343987464904785,4.34452486038208,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:48,0.6.3
tiled_geglu,liger,full,speed,ms,T,sequence length,4096,8.653023719787598,8.653023719787598,8.653023719787598,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:48,0.6.3
tiled_geglu,liger,full,speed,ms,T,sequence length,8192,16.909311294555664,16.909311294555664,16.909311294555664,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:48,0.6.3
tiled_geglu,liger,full,speed,ms,T,sequence length,16384,33.63123321533203,33.63123321533203,33.63123321533203,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:48,0.6.3
tiled_geglu,liger_tiled,full,speed,ms,T,sequence length,1024,3.353935956954956,3.353523015975952,3.35434889793396,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:49,0.6.3
tiled_geglu,liger_tiled,full,speed,ms,T,sequence length,2048,6.023168087005615,6.023168087005615,6.023168087005615,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:49,0.6.3
tiled_geglu,liger_tiled,full,speed,ms,T,sequence length,4096,11.495424270629883,11.495424270629883,11.495424270629883,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:49,0.6.3
tiled_geglu,liger_tiled,full,speed,ms,T,sequence length,8192,23.68614387512207,23.68614387512207,23.68614387512207,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:49,0.6.3
tiled_geglu,liger_tiled,full,speed,ms,T,sequence length,16384,47.478782653808594,47.478782653808594,47.478782653808594,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:49,0.6.3
tiled_geglu,liger,forward,speed,ms,T,sequence length,1024,0.6614400148391724,0.6594560146331787,0.6635519862174988,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:52,0.6.3
tiled_geglu,liger,forward,speed,ms,T,sequence length,2048,1.3471999168395996,1.346560001373291,1.3475840091705322,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:52,0.6.3
tiled_geglu,liger,forward,speed,ms,T,sequence length,4096,2.752511978149414,2.7261502742767334,2.7844607830047607,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:52,0.6.3
tiled_geglu,liger,forward,speed,ms,T,sequence length,8192,5.433343887329102,5.433343887329102,5.433343887329102,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:52,0.6.3
tiled_geglu,liger,forward,speed,ms,T,sequence length,16384,10.712063789367676,10.712063789367676,10.712063789367676,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:52,0.6.3
tiled_geglu,liger_tiled,forward,speed,ms,T,sequence length,1024,0.7403519749641418,0.7402047514915466,0.7413759827613831,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:53,0.6.3
tiled_geglu,liger_tiled,forward,speed,ms,T,sequence length,2048,1.3941760063171387,1.3895679712295532,1.398144006729126,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:53,0.6.3
tiled_geglu,liger_tiled,forward,speed,ms,T,sequence length,4096,2.7586560249328613,2.7585408687591553,2.759884834289551,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:53,0.6.3
tiled_geglu,liger_tiled,forward,speed,ms,T,sequence length,8192,5.789696216583252,5.789696216583252,5.789696216583252,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:53,0.6.3
tiled_geglu,liger_tiled,forward,speed,ms,T,sequence length,16384,11.810815811157227,11.810815811157227,11.810815811157227,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:53,0.6.3
tiled_geglu,liger,backward,speed,ms,T,sequence length,1024,1.491968035697937,1.4916608333587646,1.4940160512924194,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:56,0.6.3
tiled_geglu,liger,backward,speed,ms,T,sequence length,2048,3.0185279846191406,3.0131328105926514,3.0555264949798584,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:56,0.6.3
tiled_geglu,liger,backward,speed,ms,T,sequence length,4096,6.021120071411133,6.021120071411133,6.021120071411133,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:56,0.6.3
tiled_geglu,liger,backward,speed,ms,T,sequence length,8192,11.512767791748047,11.512767791748047,11.512767791748047,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:56,0.6.3
tiled_geglu,liger,backward,speed,ms,T,sequence length,16384,22.806528091430664,22.806528091430664,22.806528091430664,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:56,0.6.3
tiled_geglu,liger_tiled,backward,speed,ms,T,sequence length,1024,2.6060800552368164,2.6053311824798584,2.607308864593506,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:57,0.6.3
tiled_geglu,liger_tiled,backward,speed,ms,T,sequence length,2048,4.665375709533691,4.664742469787598,4.666009426116943,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:57,0.6.3
tiled_geglu,liger_tiled,backward,speed,ms,T,sequence length,4096,8.71731185913086,8.71731185913086,8.71731185913086,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:57,0.6.3
tiled_geglu,liger_tiled,backward,speed,ms,T,sequence length,8192,17.99782371520996,17.99782371520996,17.99782371520996,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:57,0.6.3
tiled_geglu,liger_tiled,backward,speed,ms,T,sequence length,16384,35.64400100708008,35.64400100708008,35.64400100708008,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:57,0.6.3
tiled_geglu,liger,full,memory,MB,T,sequence length,1024,232.25,232.25,232.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:00,0.6.3
tiled_geglu,liger,full,memory,MB,T,sequence length,2048,336.25,336.25,336.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:00,0.6.3
tiled_geglu,liger,full,memory,MB,T,sequence length,4096,544.25,544.25,544.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:00,0.6.3
tiled_geglu,liger,full,memory,MB,T,sequence length,8192,960.25,960.25,960.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:00,0.6.3
tiled_geglu,liger,full,memory,MB,T,sequence length,16384,1792.25,1792.25,1792.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:00,0.6.3
tiled_geglu,liger_tiled,full,memory,MB,T,sequence length,1024,186.25,186.25,186.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:00,0.6.3
tiled_geglu,liger_tiled,full,memory,MB,T,sequence length,2048,244.25,244.25,244.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:00,0.6.3
tiled_geglu,liger_tiled,full,memory,MB,T,sequence length,4096,360.25,360.25,360.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:00,0.6.3
tiled_geglu,liger_tiled,full,memory,MB,T,sequence length,8192,592.25,592.25,592.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:00,0.6.3
tiled_geglu,liger_tiled,full,memory,MB,T,sequence length,16384,1056.25,1056.25,1056.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:00,0.6.3
tiled_geglu,liger,forward,memory,MB,T,sequence length,1024,128.25,128.25,128.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:03,0.6.3
tiled_geglu,liger,forward,memory,MB,T,sequence length,2048,192.25,192.25,192.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:03,0.6.3
tiled_geglu,liger,forward,memory,MB,T,sequence length,4096,320.25,320.25,320.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:03,0.6.3
tiled_geglu,liger,forward,memory,MB,T,sequence length,8192,576.25,576.25,576.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:03,0.6.3
tiled_geglu,liger,forward,memory,MB,T,sequence length,16384,1088.25,1088.25,1088.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:03,0.6.3
tiled_geglu,liger_tiled,forward,memory,MB,T,sequence length,1024,92.25,92.25,92.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:04,0.6.3
tiled_geglu,liger_tiled,forward,memory,MB,T,sequence length,2048,120.25,120.25,120.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:04,0.6.3
tiled_geglu,liger_tiled,forward,memory,MB,T,sequence length,4096,176.25,176.25,176.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:04,0.6.3
tiled_geglu,liger_tiled,forward,memory,MB,T,sequence length,8192,288.25,288.25,288.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:04,0.6.3
tiled_geglu,liger_tiled,forward,memory,MB,T,sequence length,16384,512.25,512.25,512.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:04,0.6.3
tiled_geglu,liger,backward,memory,MB,T,sequence length,1024,232.25,232.25,232.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:05,0.6.3
tiled_geglu,liger,backward,memory,MB,T,sequence length,2048,336.25,336.25,336.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:05,0.6.3
tiled_geglu,liger,backward,memory,MB,T,sequence length,4096,544.25,544.25,544.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:05,0.6.3
tiled_geglu,liger,backward,memory,MB,T,sequence length,8192,960.25,960.25,960.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:05,0.6.3
tiled_geglu,liger,backward,memory,MB,T,sequence length,16384,1792.25,1792.25,1792.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:05,0.6.3
tiled_geglu,liger_tiled,backward,memory,MB,T,sequence length,1024,186.25,186.25,186.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:06,0.6.3
tiled_geglu,liger_tiled,backward,memory,MB,T,sequence length,2048,244.25,244.25,244.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:06,0.6.3
tiled_geglu,liger_tiled,backward,memory,MB,T,sequence length,4096,360.25,360.25,360.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:06,0.6.3
tiled_geglu,liger_tiled,backward,memory,MB,T,sequence length,8192,592.25,592.25,592.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:06,0.6.3
tiled_geglu,liger_tiled,backward,memory,MB,T,sequence length,16384,1056.25,1056.25,1056.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:06,0.6.3
tiled_swiglu,liger,full,speed,ms,T,sequence length,1024,2.165760040283203,2.164659261703491,2.167193651199341,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:10,0.6.3
tiled_swiglu,liger,full,speed,ms,T,sequence length,2048,4.371456146240234,4.368383884429932,4.374527931213379,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:10,0.6.3
tiled_swiglu,liger,full,speed,ms,T,sequence length,4096,8.935423851013184,8.935423851013184,8.935423851013184,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:10,0.6.3
tiled_swiglu,liger,full,speed,ms,T,sequence length,8192,17.078943252563477,17.078943252563477,17.078943252563477,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:10,0.6.3
tiled_swiglu,liger,full,speed,ms,T,sequence length,16384,33.74857711791992,33.74857711791992,33.74857711791992,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:10,0.6.3
tiled_swiglu,liger_tiled,full,speed,ms,T,sequence length,1024,3.3510398864746094,3.3507328033447266,3.3513472080230713,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:11,0.6.3
tiled_swiglu,liger_tiled,full,speed,ms,T,sequence length,2048,6.023168087005615,6.023168087005615,6.023168087005615,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:11,0.6.3
tiled_swiglu,liger_tiled,full,speed,ms,T,sequence length,4096,11.609087944030762,11.609087944030762,11.609087944030762,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:11,0.6.3
tiled_swiglu,liger_tiled,full,speed,ms,T,sequence length,8192,23.8591365814209,23.8591365814209,23.8591365814209,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:11,0.6.3
tiled_swiglu,liger_tiled,full,speed,ms,T,sequence length,16384,47.721473693847656,47.721473693847656,47.721473693847656,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:11,0.6.3
tiled_swiglu,liger,forward,speed,ms,T,sequence length,1024,0.6594560146331787,0.6594560146331787,0.6604800224304199,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:14,0.6.3
tiled_swiglu,liger,forward,speed,ms,T,sequence length,2048,1.3537280559539795,1.3527040481567383,1.3547519445419312,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:14,0.6.3
tiled_swiglu,liger,forward,speed,ms,T,sequence length,4096,2.7152960300445557,2.715123176574707,2.7155072689056396,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:14,0.6.3
tiled_swiglu,liger,forward,speed,ms,T,sequence length,8192,5.3361921310424805,5.3361921310424805,5.3361921310424805,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:14,0.6.3
tiled_swiglu,liger,forward,speed,ms,T,sequence length,16384,10.870783805847168,10.870783805847168,10.870783805847168,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:14,0.6.3
tiled_swiglu,liger_tiled,forward,speed,ms,T,sequence length,1024,0.7395360469818115,0.7383040189743042,0.7413759827613831,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:15,0.6.3
tiled_swiglu,liger_tiled,forward,speed,ms,T,sequence length,2048,1.3965599536895752,1.387935996055603,1.4024640321731567,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:15,0.6.3
tiled_swiglu,liger_tiled,forward,speed,ms,T,sequence length,4096,2.7778561115264893,2.777395248413086,2.7780096530914307,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:15,0.6.3
tiled_swiglu,liger_tiled,forward,speed,ms,T,sequence length,8192,5.829631805419922,5.829631805419922,5.829631805419922,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:15,0.6.3
tiled_swiglu,liger_tiled,forward,speed,ms,T,sequence length,16384,11.841535568237305,11.841535568237305,11.841535568237305,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:15,0.6.3
tiled_swiglu,liger,backward,speed,ms,T,sequence length,1024,1.4970879554748535,1.4961408376693726,1.4970879554748535,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:17,0.6.3
tiled_swiglu,liger,backward,speed,ms,T,sequence length,2048,3.052351951599121,3.0518529415130615,3.0550782680511475,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:17,0.6.3
tiled_swiglu,liger,backward,speed,ms,T,sequence length,4096,6.074687957763672,6.074687957763672,6.074687957763672,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:17,0.6.3
tiled_swiglu,liger,backward,speed,ms,T,sequence length,8192,11.630592346191406,11.630592346191406,11.630592346191406,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:17,0.6.3
tiled_swiglu,liger,backward,speed,ms,T,sequence length,16384,22.76793670654297,22.76793670654297,22.76793670654297,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:17,0.6.3
tiled_swiglu,liger_tiled,backward,speed,ms,T,sequence length,1024,2.6021440029144287,2.6000702381134033,2.6032767295837402,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:18,0.6.3
tiled_swiglu,liger_tiled,backward,speed,ms,T,sequence length,2048,4.641791820526123,4.641791820526123,4.641791820526123,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:18,0.6.3
tiled_swiglu,liger_tiled,backward,speed,ms,T,sequence length,4096,8.761343955993652,8.761343955993652,8.761343955993652,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:18,0.6.3
tiled_swiglu,liger_tiled,backward,speed,ms,T,sequence length,8192,17.966079711914062,17.966079711914062,17.966079711914062,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:18,0.6.3
tiled_swiglu,liger_tiled,backward,speed,ms,T,sequence length,16384,35.657344818115234,35.657344818115234,35.657344818115234,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:18,0.6.3
tiled_swiglu,liger,full,memory,MB,T,sequence length,1024,232.25,232.25,232.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:21,0.6.3
tiled_swiglu,liger,full,memory,MB,T,sequence length,2048,336.25,336.25,336.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:21,0.6.3
tiled_swiglu,liger,full,memory,MB,T,sequence length,4096,544.25,544.25,544.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:21,0.6.3
tiled_swiglu,liger,full,memory,MB,T,sequence length,8192,960.25,960.25,960.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:21,0.6.3
tiled_swiglu,liger,full,memory,MB,T,sequence length,16384,1792.25,1792.25,1792.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:21,0.6.3
tiled_swiglu,liger_tiled,full,memory,MB,T,sequence length,1024,186.25,186.25,186.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:22,0.6.3
tiled_swiglu,liger_tiled,full,memory,MB,T,sequence length,2048,244.25,244.25,244.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:22,0.6.3
tiled_swiglu,liger_tiled,full,memory,MB,T,sequence length,4096,360.25,360.25,360.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:22,0.6.3
tiled_swiglu,liger_tiled,full,memory,MB,T,sequence length,8192,592.25,592.25,592.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:22,0.6.3
tiled_swiglu,liger_tiled,full,memory,MB,T,sequence length,16384,1056.25,1056.25,1056.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:22,0.6.3
tiled_swiglu,liger,forward,memory,MB,T,sequence length,1024,128.25,128.25,128.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:25,0.6.3
tiled_swiglu,liger,forward,memory,MB,T,sequence length,2048,192.25,192.25,192.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:25,0.6.3
tiled_swiglu,liger,forward,memory,MB,T,sequence length,4096,320.25,320.25,320.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:25,0.6.3
tiled_swiglu,liger,forward,memory,MB,T,sequence length,8192,576.25,576.25,576.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:25,0.6.3
tiled_swiglu,liger,forward,memory,MB,T,sequence length,16384,1088.25,1088.25,1088.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:25,0.6.3
tiled_swiglu,liger_tiled,forward,memory,MB,T,sequence length,1024,92.25,92.25,92.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:25,0.6.3
tiled_swiglu,liger_tiled,forward,memory,MB,T,sequence length,2048,120.25,120.25,120.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:25,0.6.3
tiled_swiglu,liger_tiled,forward,memory,MB,T,sequence length,4096,176.25,176.25,176.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:25,0.6.3
tiled_swiglu,liger_tiled,forward,memory,MB,T,sequence length,8192,288.25,288.25,288.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:25,0.6.3
tiled_swiglu,liger_tiled,forward,memory,MB,T,sequence length,16384,512.25,512.25,512.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:25,0.6.3
tiled_swiglu,liger,backward,memory,MB,T,sequence length,1024,232.25,232.25,232.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:27,0.6.3
tiled_swiglu,liger,backward,memory,MB,T,sequence length,2048,336.25,336.25,336.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:27,0.6.3
tiled_swiglu,liger,backward,memory,MB,T,sequence length,4096,544.25,544.25,544.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:27,0.6.3
tiled_swiglu,liger,backward,memory,MB,T,sequence length,8192,960.25,960.25,960.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:27,0.6.3
tiled_swiglu,liger,backward,memory,MB,T,sequence length,16384,1792.25,1792.25,1792.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:27,0.6.3
tiled_swiglu,liger_tiled,backward,memory,MB,T,sequence length,1024,186.25,186.25,186.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:28,0.6.3
tiled_swiglu,liger_tiled,backward,memory,MB,T,sequence length,2048,244.25,244.25,244.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:28,0.6.3
tiled_swiglu,liger_tiled,backward,memory,MB,T,sequence length,4096,360.25,360.25,360.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:28,0.6.3
tiled_swiglu,liger_tiled,backward,memory,MB,T,sequence length,8192,592.25,592.25,592.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:28,0.6.3
tiled_swiglu,liger_tiled,backward,memory,MB,T,sequence length,16384,1056.25,1056.25,1056.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:28,0.6.3
tiled_geglu,huggingface,full,speed,ms,T,sequence length,1024,2.3357439041137695,2.3357439041137695,2.3375871181488037,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:47,0.6.3
tiled_geglu,huggingface,full,speed,ms,T,sequence length,2048,4.764671802520752,4.764671802520752,4.764671802520752,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:47,0.6.3
tiled_geglu,huggingface,full,speed,ms,T,sequence length,4096,9.4236478805542,9.4236478805542,9.4236478805542,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:47,0.6.3
tiled_geglu,huggingface,full,speed,ms,T,sequence length,8192,17.628543853759766,17.628543853759766,17.628543853759766,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:47,0.6.3
tiled_geglu,huggingface,full,speed,ms,T,sequence length,16384,35.06790542602539,35.06790542602539,35.06790542602539,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:47,0.6.3
tiled_geglu,deepspeed_tiled,full,speed,ms,T,sequence length,1024,3.418976068496704,3.4176511764526367,3.4203009605407715,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:51,0.6.3
tiled_geglu,deepspeed_tiled,full,speed,ms,T,sequence length,2048,6.158143997192383,6.158143997192383,6.158143997192383,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:51,0.6.3
tiled_geglu,deepspeed_tiled,full,speed,ms,T,sequence length,4096,11.934720039367676,11.934720039367676,11.934720039367676,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:51,0.6.3
tiled_geglu,deepspeed_tiled,full,speed,ms,T,sequence length,8192,24.731647491455078,24.731647491455078,24.731647491455078,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:51,0.6.3
tiled_geglu,deepspeed_tiled,full,speed,ms,T,sequence length,16384,49.46227264404297,49.46227264404297,49.46227264404297,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:51,0.6.3
tiled_geglu,huggingface,forward,speed,ms,T,sequence length,1024,0.6743040084838867,0.6736640334129333,0.677068829536438,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:52,0.6.3
tiled_geglu,huggingface,forward,speed,ms,T,sequence length,2048,1.418239951133728,1.418239951133728,1.421120047569275,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:52,0.6.3
tiled_geglu,huggingface,forward,speed,ms,T,sequence length,4096,2.88972806930542,2.889113664627075,2.8909568786621094,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:52,0.6.3
tiled_geglu,huggingface,forward,speed,ms,T,sequence length,8192,5.701375961303711,5.701375961303711,5.701375961303711,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:52,0.6.3
tiled_geglu,huggingface,forward,speed,ms,T,sequence length,16384,11.276288032531738,11.276288032531738,11.276288032531738,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:52,0.6.3
tiled_geglu,deepspeed_tiled,forward,speed,ms,T,sequence length,1024,0.7433919906616211,0.7423999905586243,0.7444480061531067,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:54,0.6.3
tiled_geglu,deepspeed_tiled,forward,speed,ms,T,sequence length,2048,1.4137760400772095,1.4131200313568115,1.4152319431304932,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:54,0.6.3
tiled_geglu,deepspeed_tiled,forward,speed,ms,T,sequence length,4096,2.8241920471191406,2.823500871658325,2.8266496658325195,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:54,0.6.3
tiled_geglu,deepspeed_tiled,forward,speed,ms,T,sequence length,8192,6.087679862976074,6.087679862976074,6.087679862976074,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:54,0.6.3
tiled_geglu,deepspeed_tiled,forward,speed,ms,T,sequence length,16384,12.353535652160645,12.353535652160645,12.353535652160645,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:54,0.6.3
tiled_geglu,huggingface,backward,speed,ms,T,sequence length,1024,1.5499199628829956,1.5489535331726074,1.5523840188980103,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:55,0.6.3
tiled_geglu,huggingface,backward,speed,ms,T,sequence length,2048,3.171328067779541,3.169484853744507,3.173171281814575,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:55,0.6.3
tiled_geglu,huggingface,backward,speed,ms,T,sequence length,4096,6.263807773590088,6.263807773590088,6.263807773590088,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:55,0.6.3
tiled_geglu,huggingface,backward,speed,ms,T,sequence length,8192,12.046143531799316,12.046143531799316,12.046143531799316,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:55,0.6.3
tiled_geglu,huggingface,backward,speed,ms,T,sequence length,16384,23.839744567871094,23.839744567871094,23.839744567871094,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:55,0.6.3
tiled_geglu,deepspeed_tiled,backward,speed,ms,T,sequence length,1024,2.6757121086120605,2.6755776405334473,2.676710367202759,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:58,0.6.3
tiled_geglu,deepspeed_tiled,backward,speed,ms,T,sequence length,2048,4.7329277992248535,4.7329277992248535,4.7329277992248535,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:58,0.6.3
tiled_geglu,deepspeed_tiled,backward,speed,ms,T,sequence length,4096,9.078783988952637,9.078783988952637,9.078783988952637,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:58,0.6.3
tiled_geglu,deepspeed_tiled,backward,speed,ms,T,sequence length,8192,18.63680076599121,18.63680076599121,18.63680076599121,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:58,0.6.3
tiled_geglu,deepspeed_tiled,backward,speed,ms,T,sequence length,16384,37.06163024902344,37.06163024902344,37.06163024902344,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:58,0.6.3
tiled_geglu,huggingface,full,memory,MB,T,sequence length,1024,264.25,264.25,264.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:59,0.6.3
tiled_geglu,huggingface,full,memory,MB,T,sequence length,2048,400.25,400.25,400.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:59,0.6.3
tiled_geglu,huggingface,full,memory,MB,T,sequence length,4096,688.25,688.25,688.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:59,0.6.3
tiled_geglu,huggingface,full,memory,MB,T,sequence length,8192,1264.25,1264.25,1264.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:59,0.6.3
tiled_geglu,huggingface,full,memory,MB,T,sequence length,16384,2416.25,2416.25,2416.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:22:59,0.6.3
tiled_geglu,deepspeed_tiled,full,memory,MB,T,sequence length,1024,190.25,190.25,190.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:02,0.6.3
tiled_geglu,deepspeed_tiled,full,memory,MB,T,sequence length,2048,252.25,252.25,252.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:02,0.6.3
tiled_geglu,deepspeed_tiled,full,memory,MB,T,sequence length,4096,376.25,376.25,376.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:02,0.6.3
tiled_geglu,deepspeed_tiled,full,memory,MB,T,sequence length,8192,640.25,640.25,640.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:02,0.6.3
tiled_geglu,deepspeed_tiled,full,memory,MB,T,sequence length,16384,1168.25,1168.25,1168.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:02,0.6.3
tiled_geglu,huggingface,forward,memory,MB,T,sequence length,1024,144.25,144.25,144.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:02,0.6.3
tiled_geglu,huggingface,forward,memory,MB,T,sequence length,2048,224.25,224.25,224.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:02,0.6.3
tiled_geglu,huggingface,forward,memory,MB,T,sequence length,4096,384.25,384.25,384.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:02,0.6.3
tiled_geglu,huggingface,forward,memory,MB,T,sequence length,8192,704.25,704.25,704.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:02,0.6.3
tiled_geglu,huggingface,forward,memory,MB,T,sequence length,16384,1344.25,1344.25,1344.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:02,0.6.3
tiled_geglu,deepspeed_tiled,forward,memory,MB,T,sequence length,1024,90.25,90.25,90.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:04,0.6.3
tiled_geglu,deepspeed_tiled,forward,memory,MB,T,sequence length,2048,116.25,116.25,116.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:04,0.6.3
tiled_geglu,deepspeed_tiled,forward,memory,MB,T,sequence length,4096,168.25,168.25,168.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:04,0.6.3
tiled_geglu,deepspeed_tiled,forward,memory,MB,T,sequence length,8192,272.25,272.25,272.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:04,0.6.3
tiled_geglu,deepspeed_tiled,forward,memory,MB,T,sequence length,16384,480.25,480.25,480.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:04,0.6.3
tiled_geglu,huggingface,backward,memory,MB,T,sequence length,1024,264.25,264.25,264.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:05,0.6.3
tiled_geglu,huggingface,backward,memory,MB,T,sequence length,2048,400.25,400.25,400.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:05,0.6.3
tiled_geglu,huggingface,backward,memory,MB,T,sequence length,4096,688.25,688.25,688.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:05,0.6.3
tiled_geglu,huggingface,backward,memory,MB,T,sequence length,8192,1264.25,1264.25,1264.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:05,0.6.3
tiled_geglu,huggingface,backward,memory,MB,T,sequence length,16384,2416.25,2416.25,2416.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:05,0.6.3
tiled_geglu,deepspeed_tiled,backward,memory,MB,T,sequence length,1024,190.25,190.25,190.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:07,0.6.3
tiled_geglu,deepspeed_tiled,backward,memory,MB,T,sequence length,2048,252.25,252.25,252.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:07,0.6.3
tiled_geglu,deepspeed_tiled,backward,memory,MB,T,sequence length,4096,376.25,376.25,376.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:07,0.6.3
tiled_geglu,deepspeed_tiled,backward,memory,MB,T,sequence length,8192,640.25,640.25,640.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:07,0.6.3
tiled_geglu,deepspeed_tiled,backward,memory,MB,T,sequence length,16384,1168.25,1168.25,1168.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""gelu_pytorch_tanh"", ""activation_type"": ""geglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:07,0.6.3
tiled_swiglu,huggingface,full,speed,ms,T,sequence length,1024,2.2517759799957275,2.2517759799957275,2.254848003387451,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:08,0.6.3
tiled_swiglu,huggingface,full,speed,ms,T,sequence length,2048,4.588511943817139,4.587302207946777,4.5897216796875,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:08,0.6.3
tiled_swiglu,huggingface,full,speed,ms,T,sequence length,4096,9.233407974243164,9.233407974243164,9.233407974243164,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:08,0.6.3
tiled_swiglu,huggingface,full,speed,ms,T,sequence length,8192,17.869823455810547,17.869823455810547,17.869823455810547,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:08,0.6.3
tiled_swiglu,huggingface,full,speed,ms,T,sequence length,16384,35.34422302246094,35.34422302246094,35.34422302246094,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:08,0.6.3
tiled_swiglu,deepspeed_tiled,full,speed,ms,T,sequence length,1024,3.4257922172546387,3.424870491027832,3.426713705062866,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:12,0.6.3
tiled_swiglu,deepspeed_tiled,full,speed,ms,T,sequence length,2048,6.155263900756836,6.155263900756836,6.155263900756836,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:12,0.6.3
tiled_swiglu,deepspeed_tiled,full,speed,ms,T,sequence length,4096,11.92959976196289,11.92959976196289,11.92959976196289,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:12,0.6.3
tiled_swiglu,deepspeed_tiled,full,speed,ms,T,sequence length,8192,24.815616607666016,24.815616607666016,24.815616607666016,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:12,0.6.3
tiled_swiglu,deepspeed_tiled,full,speed,ms,T,sequence length,16384,49.62918472290039,49.62918472290039,49.62918472290039,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:12,0.6.3
tiled_swiglu,huggingface,forward,speed,ms,T,sequence length,1024,0.6748160123825073,0.6737920045852661,0.6758400201797485,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:13,0.6.3
tiled_swiglu,huggingface,forward,speed,ms,T,sequence length,2048,1.4332799911499023,1.4325759410858154,1.4335999488830566,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:13,0.6.3
tiled_swiglu,huggingface,forward,speed,ms,T,sequence length,4096,2.91212797164917,2.904217481613159,2.9146623611450195,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:13,0.6.3
tiled_swiglu,huggingface,forward,speed,ms,T,sequence length,8192,5.658976078033447,5.658976078033447,5.658976078033447,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:13,0.6.3
tiled_swiglu,huggingface,forward,speed,ms,T,sequence length,16384,11.341952323913574,11.341952323913574,11.341952323913574,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:13,0.6.3
tiled_swiglu,deepspeed_tiled,forward,speed,ms,T,sequence length,1024,0.7454720139503479,0.7429631948471069,0.7456768155097961,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:15,0.6.3
tiled_swiglu,deepspeed_tiled,forward,speed,ms,T,sequence length,2048,1.4120960235595703,1.410048007965088,1.4120960235595703,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:15,0.6.3
tiled_swiglu,deepspeed_tiled,forward,speed,ms,T,sequence length,4096,2.825216054916382,2.825216054916382,2.8264448642730713,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:15,0.6.3
tiled_swiglu,deepspeed_tiled,forward,speed,ms,T,sequence length,8192,6.077439785003662,6.077439785003662,6.077439785003662,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:15,0.6.3
tiled_swiglu,deepspeed_tiled,forward,speed,ms,T,sequence length,16384,12.356608390808105,12.356608390808105,12.356608390808105,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:15,0.6.3
tiled_swiglu,huggingface,backward,speed,ms,T,sequence length,1024,1.551360011100769,1.5511807203292847,1.5532032251358032,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:16,0.6.3
tiled_swiglu,huggingface,backward,speed,ms,T,sequence length,2048,3.1928319931030273,3.1885311603546143,3.1971328258514404,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:16,0.6.3
tiled_swiglu,huggingface,backward,speed,ms,T,sequence length,4096,6.273248195648193,6.273248195648193,6.273248195648193,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:16,0.6.3
tiled_swiglu,huggingface,backward,speed,ms,T,sequence length,8192,12.058752059936523,12.058752059936523,12.058752059936523,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:16,0.6.3
tiled_swiglu,huggingface,backward,speed,ms,T,sequence length,16384,23.853055953979492,23.853055953979492,23.853055953979492,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:16,0.6.3
tiled_swiglu,deepspeed_tiled,backward,speed,ms,T,sequence length,1024,2.6746881008148193,2.6728639602661133,2.6789886951446533,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:20,0.6.3
tiled_swiglu,deepspeed_tiled,backward,speed,ms,T,sequence length,2048,4.739071846008301,4.739071846008301,4.739071846008301,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:20,0.6.3
tiled_swiglu,deepspeed_tiled,backward,speed,ms,T,sequence length,4096,9.084927558898926,9.084927558898926,9.084927558898926,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:20,0.6.3
tiled_swiglu,deepspeed_tiled,backward,speed,ms,T,sequence length,8192,18.729759216308594,18.729759216308594,18.729759216308594,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:20,0.6.3
tiled_swiglu,deepspeed_tiled,backward,speed,ms,T,sequence length,16384,37.13724899291992,37.13724899291992,37.13724899291992,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:20,0.6.3
tiled_swiglu,huggingface,full,memory,MB,T,sequence length,1024,264.25,264.25,264.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:20,0.6.3
tiled_swiglu,huggingface,full,memory,MB,T,sequence length,2048,400.25,400.25,400.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:20,0.6.3
tiled_swiglu,huggingface,full,memory,MB,T,sequence length,4096,688.25,688.25,688.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:20,0.6.3
tiled_swiglu,huggingface,full,memory,MB,T,sequence length,8192,1264.25,1264.25,1264.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:20,0.6.3
tiled_swiglu,huggingface,full,memory,MB,T,sequence length,16384,2416.25,2416.25,2416.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:20,0.6.3
tiled_swiglu,deepspeed_tiled,full,memory,MB,T,sequence length,1024,190.25,190.25,190.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:23,0.6.3
tiled_swiglu,deepspeed_tiled,full,memory,MB,T,sequence length,2048,252.25,252.25,252.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:23,0.6.3
tiled_swiglu,deepspeed_tiled,full,memory,MB,T,sequence length,4096,376.25,376.25,376.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:23,0.6.3
tiled_swiglu,deepspeed_tiled,full,memory,MB,T,sequence length,8192,640.25,640.25,640.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:23,0.6.3
tiled_swiglu,deepspeed_tiled,full,memory,MB,T,sequence length,16384,1168.25,1168.25,1168.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:23,0.6.3
tiled_swiglu,huggingface,forward,memory,MB,T,sequence length,1024,144.25,144.25,144.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:24,0.6.3
tiled_swiglu,huggingface,forward,memory,MB,T,sequence length,2048,224.25,224.25,224.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:24,0.6.3
tiled_swiglu,huggingface,forward,memory,MB,T,sequence length,4096,384.25,384.25,384.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:24,0.6.3
tiled_swiglu,huggingface,forward,memory,MB,T,sequence length,8192,704.25,704.25,704.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:24,0.6.3
tiled_swiglu,huggingface,forward,memory,MB,T,sequence length,16384,1344.25,1344.25,1344.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:24,0.6.3
tiled_swiglu,deepspeed_tiled,forward,memory,MB,T,sequence length,1024,90.25,90.25,90.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:26,0.6.3
tiled_swiglu,deepspeed_tiled,forward,memory,MB,T,sequence length,2048,116.25,116.25,116.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:26,0.6.3
tiled_swiglu,deepspeed_tiled,forward,memory,MB,T,sequence length,4096,168.25,168.25,168.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:26,0.6.3
tiled_swiglu,deepspeed_tiled,forward,memory,MB,T,sequence length,8192,272.25,272.25,272.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:26,0.6.3
tiled_swiglu,deepspeed_tiled,forward,memory,MB,T,sequence length,16384,480.25,480.25,480.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:26,0.6.3
tiled_swiglu,huggingface,backward,memory,MB,T,sequence length,1024,264.25,264.25,264.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:26,0.6.3
tiled_swiglu,huggingface,backward,memory,MB,T,sequence length,2048,400.25,400.25,400.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:26,0.6.3
tiled_swiglu,huggingface,backward,memory,MB,T,sequence length,4096,688.25,688.25,688.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:26,0.6.3
tiled_swiglu,huggingface,backward,memory,MB,T,sequence length,8192,1264.25,1264.25,1264.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:26,0.6.3
tiled_swiglu,huggingface,backward,memory,MB,T,sequence length,16384,2416.25,2416.25,2416.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:26,0.6.3
tiled_swiglu,deepspeed_tiled,backward,memory,MB,T,sequence length,1024,190.25,190.25,190.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:29,0.6.3
tiled_swiglu,deepspeed_tiled,backward,memory,MB,T,sequence length,2048,252.25,252.25,252.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:29,0.6.3
tiled_swiglu,deepspeed_tiled,backward,memory,MB,T,sequence length,4096,376.25,376.25,376.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:29,0.6.3
tiled_swiglu,deepspeed_tiled,backward,memory,MB,T,sequence length,8192,640.25,640.25,640.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:29,0.6.3
tiled_swiglu,deepspeed_tiled,backward,memory,MB,T,sequence length,16384,1168.25,1168.25,1168.25,"{""bsz"": 2, ""hidden_size"": 2048, ""intermediate_size"": 4096, ""hidden_act"": ""silu"", ""activation_type"": ""swiglu"", ""num_shards"": 4, ""dtype"": ""torch.bfloat16""}",NVIDIA GeForce RTX 4090,2025-11-11 06:23:29,0.6.3
Can't render this file because it is too large.
+397
View File
@@ -0,0 +1,397 @@
import math
import torch
import torch.nn as nn
import triton
from transformers.models.llama.configuration_llama import LlamaConfig
from transformers.models.llama.modeling_llama import LlamaMLP
from utils import QUANTILES
from utils import SingleBenchmarkRunInput
from utils import SingleBenchmarkRunOutput
from utils import _test_memory
from utils import parse_benchmark_script_args
from utils import run_benchmarks
from liger_kernel.transformers.geglu import LigerGEGLUMLP
from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
from liger_kernel.transformers.tiled_mlp import LigerTiledGEGLUMLP
from liger_kernel.transformers.tiled_mlp import LigerTiledSwiGLUMLP
from liger_kernel.utils import infer_device
device = infer_device()
# DeepSpeed TiledMLP implementation
# Based on: https://github.com/deepspeedai/DeepSpeed/blob/v0.18.2/deepspeed/runtime/sequence_parallel/ulysses_sp.py#L838
class DeepSpeedTiledMLP(torch.autograd.Function):
"""
DeepSpeed's TiledMLP implementation for fair comparison.
This is the actual DeepSpeed algorithm that performs tiled MLP computation
to massively reduce memory usage with very long sequence lengths.
This module re-computes forward in the backward, so forward occurs twice per iteration.
"""
@staticmethod
def forward(ctx, fn, self, x, shards, compute_params) -> torch.Tensor:
ctx.fn = fn
ctx.self = self
ctx.shards = shards
ctx.compute_params = [p for p in compute_params if p.requires_grad] if compute_params else []
ctx.save_for_backward(x)
# x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size] (moe experts)
x_shards = list(torch.chunk(x, chunks=shards, dim=-2))
with torch.no_grad():
output_shards = [fn(self, x_shard) for x_shard in x_shards]
output_unsharded = torch.cat(output_shards, dim=-2)
return output_unsharded
@staticmethod
def backward(ctx, *grads):
fn = ctx.fn
(x,) = ctx.saved_tensors
self = ctx.self
shards = ctx.shards
compute_params = ctx.compute_params
x_requires_grad = x.requires_grad
x = x.detach()
# detach() unsets x.requires_grad, so restore it
x.requires_grad_(x_requires_grad)
# x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size] (moe experts)
hidden_size = x.shape[-1]
x_shape_orig = x.shape
# flatten bs+seqlen to avoid having stride issues when narrowing into seqlen w/ bs>1
x = x.view(-1, hidden_size)
incoming_grad = grads[0].view(-1, hidden_size)
x_grad = torch.zeros_like(x)
x_shards = list(torch.chunk(x, chunks=shards, dim=0))
for i, x_shard in enumerate(x_shards):
# Tell deepspeed not to add a new grad to its ipg bucket until the last shard is run
# XXX: DDP, FSDP will need something similar to make it work
if compute_params:
if i + 1 < shards:
for param in compute_params:
if hasattr(param, "ds_grad_is_ready"):
param.ds_grad_is_ready = False
else:
# last shard, can add the grad
for param in compute_params:
if hasattr(param, "ds_grad_is_ready"):
param.ds_grad_is_ready = True
x_shard.requires_grad_(x_requires_grad)
# if seqlen is not exactly divisible by shards the last step will be shorter than shard_step
shard_step = x_shards[i].shape[0]
shard_offset = i * x_shards[0].shape[0]
x_shard.grad = x_grad.narrow(0, shard_offset, shard_step).view_as(x_shard)
incoming_grad_shard = incoming_grad.narrow(0, shard_offset, shard_step).view_as(x_shard)
with torch.enable_grad():
output = fn(self, x_shard)
torch.autograd.backward(output, incoming_grad_shard)
# unflatten
x_grad = x_grad.view(x_shape_orig)
return (None, None, x_grad, None, None)
# DeepSpeed TiledMLP wrapper to match our interface
class DeepSpeedTiledMLPWrapper(nn.Module):
"""
Wrapper for DeepSpeed's TiledMLP to match the interface used in benchmarks.
Uses the DeepSpeed TiledMLP algorithm for memory-efficient MLP computation.
"""
def __init__(self, config, num_shards=None):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size
self.intermediate_size = config.intermediate_size
self.num_shards = num_shards
self.mlp = LlamaMLP(config=config)
def forward(self, x):
# Calculate num_shards if not provided
num_shards = self.num_shards
if num_shards is None:
hidden_size = x.shape[-1]
seqlen = x.shape[-2]
num_shards = math.ceil(seqlen / hidden_size)
num_shards = max(1, num_shards)
# Collect compute parameters for DeepSpeed ZeRO compatibility
compute_params = [
self.mlp.down_proj.weight,
self.mlp.gate_proj.weight,
self.mlp.up_proj.weight,
]
# Define the MLP forward function for DeepSpeed TiledMLP
def mlp_forward(mlp_module, x_input):
return mlp_module.down_proj(mlp_module.act_fn(mlp_module.gate_proj(x_input)) * mlp_module.up_proj(x_input))
# Use DeepSpeed's TiledMLP implementation
return DeepSpeedTiledMLP.apply(
mlp_forward,
self.mlp,
x,
num_shards,
compute_params,
)
def bench_speed_tiled_mlp(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
seq_len = input.x
bsz = input.extra_benchmark_config["bsz"]
hidden_size = input.extra_benchmark_config["hidden_size"]
intermediate_size = input.extra_benchmark_config["intermediate_size"]
hidden_act = input.extra_benchmark_config["hidden_act"]
dtype = input.extra_benchmark_config["dtype"]
num_shards = input.extra_benchmark_config.get("num_shards", None)
activation_type = input.extra_benchmark_config["activation_type"]
provider = input.kernel_provider
mode = input.kernel_operation_mode
llama_config = LlamaConfig(
hidden_size=hidden_size,
intermediate_size=intermediate_size,
hidden_act=hidden_act,
)
x_shape = (bsz, seq_len, hidden_size)
# initialize input
x = torch.randn(*x_shape, device=device, dtype=dtype, requires_grad=True)
if activation_type == "geglu":
if provider == "huggingface":
layer = LlamaMLP(config=llama_config).to(device).to(dtype)
elif provider == "liger":
layer = LigerGEGLUMLP(config=llama_config).to(device).to(dtype)
elif provider == "liger_tiled":
layer = LigerTiledGEGLUMLP(config=llama_config, num_shards=num_shards).to(device).to(dtype)
elif provider == "deepspeed_tiled":
layer = DeepSpeedTiledMLPWrapper(config=llama_config, num_shards=num_shards).to(device).to(dtype)
else:
raise ValueError(f"Invalid provider: {provider} for GEGLU")
elif activation_type == "swiglu":
if provider == "huggingface":
layer = LlamaMLP(config=llama_config).to(device).to(dtype)
elif provider == "liger":
layer = LigerSwiGLUMLP(config=llama_config).to(device).to(dtype)
elif provider == "liger_tiled":
layer = LigerTiledSwiGLUMLP(config=llama_config, num_shards=num_shards).to(device).to(dtype)
elif provider == "deepspeed_tiled":
layer = DeepSpeedTiledMLPWrapper(config=llama_config, num_shards=num_shards).to(device).to(dtype)
else:
raise ValueError(f"Invalid provider: {provider} for SwiGLU")
else:
raise ValueError(f"Invalid activation_type: {activation_type}")
def fwd():
return layer(x)
if mode == "forward":
ms_50, ms_20, ms_80 = triton.testing.do_bench(
fwd,
grad_to_none=[x],
rep=10,
quantiles=QUANTILES,
)
elif mode == "backward":
do = torch.randn_like(x)
y = fwd()
ms_50, ms_20, ms_80 = triton.testing.do_bench(
lambda: y.backward(do, retain_graph=True),
grad_to_none=[x],
rep=10,
quantiles=QUANTILES,
)
else:
def full():
y = fwd()
y.backward(torch.randn_like(y), retain_graph=True)
ms_50, ms_20, ms_80 = triton.testing.do_bench(
full,
grad_to_none=[x],
rep=10,
quantiles=QUANTILES,
)
return SingleBenchmarkRunOutput(
y_20=ms_20,
y_50=ms_50,
y_80=ms_80,
)
def bench_memory_tiled_mlp(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
seq_len = input.x
bsz = input.extra_benchmark_config["bsz"]
hidden_size = input.extra_benchmark_config["hidden_size"]
intermediate_size = input.extra_benchmark_config["intermediate_size"]
hidden_act = input.extra_benchmark_config["hidden_act"]
dtype = input.extra_benchmark_config["dtype"]
num_shards = input.extra_benchmark_config.get("num_shards", None)
activation_type = input.extra_benchmark_config["activation_type"]
provider = input.kernel_provider
mode = input.kernel_operation_mode
llama_config = LlamaConfig(
hidden_size=hidden_size,
intermediate_size=intermediate_size,
hidden_act=hidden_act,
)
x_shape = (bsz, seq_len, hidden_size)
# initialize input
x = torch.randn(*x_shape, device=device, dtype=dtype, requires_grad=True)
if activation_type == "geglu":
if provider == "huggingface":
layer = LlamaMLP(config=llama_config).to(device).to(dtype)
elif provider == "liger":
layer = LigerGEGLUMLP(config=llama_config).to(device).to(dtype)
elif provider == "liger_tiled":
layer = LigerTiledGEGLUMLP(config=llama_config, num_shards=num_shards).to(device).to(dtype)
elif provider == "deepspeed_tiled":
layer = DeepSpeedTiledMLPWrapper(config=llama_config, num_shards=num_shards).to(device).to(dtype)
else:
raise ValueError(f"Invalid provider: {provider} for GEGLU")
elif activation_type == "swiglu":
if provider == "huggingface":
layer = LlamaMLP(config=llama_config).to(device).to(dtype)
elif provider == "liger":
layer = LigerSwiGLUMLP(config=llama_config).to(device).to(dtype)
elif provider == "liger_tiled":
layer = LigerTiledSwiGLUMLP(config=llama_config, num_shards=num_shards).to(device).to(dtype)
elif provider == "deepspeed_tiled":
layer = DeepSpeedTiledMLPWrapper(config=llama_config, num_shards=num_shards).to(device).to(dtype)
else:
raise ValueError(f"Invalid provider: {provider} for SwiGLU")
else:
raise ValueError(f"Invalid activation_type: {activation_type}")
def fwd():
return layer(x)
def full():
y = fwd()
y.backward(torch.randn_like(y), retain_graph=True)
if mode == "forward":
mem_50, mem_20, mem_80 = _test_memory(
fwd,
quantiles=QUANTILES,
)
elif mode == "backward":
do = torch.randn_like(x)
y = fwd()
mem_50, mem_20, mem_80 = _test_memory(
lambda: y.backward(do, retain_graph=True),
quantiles=QUANTILES,
)
else:
mem_50, mem_20, mem_80 = _test_memory(
full,
quantiles=QUANTILES,
)
return SingleBenchmarkRunOutput(
y_20=mem_20,
y_50=mem_50,
y_80=mem_80,
)
if __name__ == "__main__":
args = parse_benchmark_script_args()
# Benchmark GEGLU variants
kernel_providers_geglu = ["huggingface", "liger", "liger_tiled", "deepspeed_tiled"]
common_configs_geglu = {
"kernel_name": "tiled_geglu",
"x_name": "T",
"x_label": "sequence length",
"x_values": [2**i for i in range(10, 15)], # 1024 to 16384
"kernel_providers": kernel_providers_geglu,
"extra_benchmark_configs": [
{
"bsz": 2,
"hidden_size": 2048,
"intermediate_size": 4096,
"hidden_act": "gelu_pytorch_tanh",
"activation_type": "geglu",
"num_shards": 4,
"dtype": torch.bfloat16,
}
],
"overwrite": args.overwrite,
}
run_benchmarks(
bench_test_fn=bench_speed_tiled_mlp,
kernel_operation_modes=["full", "forward", "backward"],
metric_name="speed",
metric_unit="ms",
**common_configs_geglu,
)
run_benchmarks(
bench_test_fn=bench_memory_tiled_mlp,
kernel_operation_modes=["full", "forward", "backward"],
metric_name="memory",
metric_unit="MB",
**common_configs_geglu,
)
# Benchmark SwiGLU variants
kernel_providers_swiglu = ["huggingface", "liger", "liger_tiled", "deepspeed_tiled"]
common_configs_swiglu = {
"kernel_name": "tiled_swiglu",
"x_name": "T",
"x_label": "sequence length",
"x_values": [2**i for i in range(10, 15)], # 1024 to 16384
"kernel_providers": kernel_providers_swiglu,
"extra_benchmark_configs": [
{
"bsz": 2,
"hidden_size": 2048,
"intermediate_size": 4096,
"hidden_act": "silu",
"activation_type": "swiglu",
"num_shards": 4,
"dtype": torch.bfloat16,
}
],
"overwrite": args.overwrite,
}
run_benchmarks(
bench_test_fn=bench_speed_tiled_mlp,
kernel_operation_modes=["full", "forward", "backward"],
metric_name="speed",
metric_unit="ms",
**common_configs_swiglu,
)
run_benchmarks(
bench_test_fn=bench_memory_tiled_mlp,
kernel_operation_modes=["full", "forward", "backward"],
metric_name="memory",
metric_unit="MB",
**common_configs_swiglu,
)
+1 -1
View File
@@ -14,7 +14,7 @@ app = modal.App("liger_tests", image=image)
repo = image.add_local_dir(ROOT_PATH, remote_path=REMOTE_ROOT_PATH)
@app.function(gpu="H100!", image=repo, timeout=60 * 60)
@app.function(gpu="H100!", image=repo, timeout=90 * 60)
def liger_tests():
import subprocess
+1 -1
View File
@@ -14,7 +14,7 @@ app = modal.App("liger_tests_bwd", image=image)
repo = image.add_local_dir(ROOT_PATH, remote_path=REMOTE_ROOT_PATH)
@app.function(gpu="H100!", image=repo, timeout=60 * 60)
@app.function(gpu="H100!", image=repo, timeout=90 * 60)
def liger_bwd_tests():
import subprocess
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "liger_kernel"
version = "0.6.3"
version = "0.6.4"
description = "Efficient Triton kernels for LLM Training"
urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
readme = { file = "README.md", content-type = "text/markdown" }
@@ -32,7 +32,7 @@ class LigerFusedLinearPPOBase(torch.autograd.Function):
epsilon_low=0.2,
epsilon_high=0.2,
beta=0.04,
loss_type="bnpo",
loss_type="dapo",
max_completion_length=None,
importance_sampling_level="token",
temperature=1.0,
@@ -60,7 +60,7 @@ class LigerFusedLinearPPOBase(torch.autograd.Function):
epsilon_low: Lower bound for clipping the importance sampling ratio
epsilon_high: Upper bound for clipping the importance sampling ratio
beta: Weight for the KL penalty
loss_type: Type of loss calculation ("grpo", "bnpo", "dr_grpo")
loss_type: Type of loss calculation ("grpo", "bnpo", "dr_grpo", "dapo")
max_completion_length: Maximum completion length required for "dr_grpo"
temperature: Temperature for the logits
compiled: Whether to use torch compile
@@ -244,6 +244,21 @@ class LigerFusedLinearPPOBase(torch.autograd.Function):
return loss_acc, tuple(final_metrics)
@staticmethod
def _compute_dapo_normalizer(attention_mask):
"""Global active tokens averaged per process."""
normalizer = attention_mask.to(torch.float32).sum()
world_size = 1
if torch.distributed.is_available() and torch.distributed.is_initialized():
import torch.distributed as dist
normalizer = normalizer.clone()
dist.all_reduce(normalizer, op=dist.ReduceOp.SUM)
world_size = dist.get_world_size()
normalizer = normalizer / world_size
return torch.clamp(normalizer, min=1.0)
@staticmethod
def _compute_chunk_loss(
input_chunk,
@@ -261,7 +276,7 @@ class LigerFusedLinearPPOBase(torch.autograd.Function):
epsilon_low=0.2,
epsilon_high=0.2,
beta=0.04,
loss_type="bnpo",
loss_type="dapo",
max_completion_length=None,
importance_sampling_level="token",
temperature=1.0,
@@ -341,10 +356,11 @@ class LigerFusedLinearPPOBase(torch.autograd.Function):
None, # grad_epsilon_low
None, # grad_epsilon_high
None, # grad_beta
None, # grad_loss_type
None, # grad_max_completion_length
None, # grad_importance_sampling_level
None, # grad_temperature
None, # grad_compiled
None, # grad_use_ref_model
None, # grad_chunk_size
None, # grad_loss_type
None, # grad_max_completion_length
)
+8 -5
View File
@@ -29,7 +29,7 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearPPOBase):
epsilon_low=0.2,
epsilon_high=0.2,
beta=0.04,
loss_type="bnpo", # ["grpo", "bnpo", "dr_grpo"]
loss_type="dapo", # ["grpo", "bnpo", "dr_grpo", "dapo"]
max_completion_length=None, # Required for dr_grpo
importance_sampling_level="token", # ["token", "sequence"] - new parameter for GSPO
**kwargs,
@@ -94,6 +94,9 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearPPOBase):
if max_completion_length is None:
raise ValueError("max_completion_length must be provided for loss_type 'dr_grpo'")
loss = (per_token_loss * attention_mask).sum() / (full_attention_mask.shape[0] * max_completion_length)
elif loss_type == "dapo":
loss_normalizer = LigerFusedLinearPPOBase._compute_dapo_normalizer(full_attention_mask)
loss = (per_token_loss * attention_mask).sum() / loss_normalizer
else:
raise ValueError(f"Unknown loss type: {loss_type}")
@@ -135,7 +138,7 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearPPOBase):
beta=0.04,
epsilon_low=0.2,
epsilon_high=0.2,
loss_type="bnpo",
loss_type="dapo",
max_completion_length=None,
importance_sampling_level="token",
temperature=1.0,
@@ -157,7 +160,7 @@ class LigerFusedLinearGRPOFunction(LigerFusedLinearPPOBase):
ref_weight (torch.Tensor, optional): Reference model weight tensor. Shape: (vocab_size, hidden_size)
ref_bias (torch.Tensor, optional): Reference model bias tensor. Shape: (vocab_size,)
beta (float): Weight for the KL penalty
loss_type (str): Type of loss calculation ("grpo", "bnpo", "dr_grpo"). Defaults to "bnpo".
loss_type (str): Type of loss calculation ("grpo", "bnpo", "dr_grpo", "dapo"). Defaults to "dapo".
max_completion_length (int, optional): Maximum completion length, required for "dr_grpo". Defaults to None.
importance_sampling_level (str): Level of importance sampling ("token" or "sequence"). Defaults to "token".
temperature (float): Temperature for the logits
@@ -235,7 +238,7 @@ class LigerFusedLinearGRPOLoss(torch.nn.Module):
chunk_size: int = 1,
epsilon_low: float = 0.2,
epsilon_high: float = 0.2,
loss_type: str = "bnpo",
loss_type: str = "dapo",
max_completion_length: Optional[int] = None,
importance_sampling_level: str = "token",
temperature: float = 1.0,
@@ -248,7 +251,7 @@ class LigerFusedLinearGRPOLoss(torch.nn.Module):
chunk_size (int): Size of chunks for processing.
epsilon_low (float): Lower bound for the importance sampling ratio.
epsilon_high (float): Upper bound for the importance sampling ratio.
loss_type (str): Type of loss calculation ("grpo", "bnpo", "dr_grpo"). Defaults to "bnpo".
loss_type (str): Type of loss calculation ("grpo", "bnpo", "dr_grpo", "dapo"). Defaults to "dapo".
max_completion_length (int, optional): Maximum completion length, required for "dr_grpo". Defaults to None.
importance_sampling_level (str): Level of importance sampling ("token" or "sequence"). Defaults to "token".
temperature (float): Temperature for the logits.
+3 -1
View File
@@ -128,7 +128,9 @@ def _grpo_loss_fwd_kernel(
per_token_loss1 = coef_1 * advantage
per_token_loss2 = coef_2 * advantage
per_token_loss = -tl.minimum(per_token_loss1, per_token_loss2)
is_clipped = per_token_loss1 < per_token_loss2
is_low_clipped = (coef_1 < 1 - EPS_LOW) & (advantage < 0)
is_high_clipped = (coef_1 > 1 + EPS_HIGH) & (advantage > 0)
is_clipped = is_low_clipped | is_high_clipped
if BETA != 0.0:
REF_LOGP += off_b * L + off_l
+81 -62
View File
@@ -1,3 +1,4 @@
import math
import operator
import torch
@@ -85,68 +86,87 @@ def _layer_norm_forward_kernel(
@triton.jit
def _layer_norm_backward_kernel(
X_ptr, # pointer to input, shape (n_rows, n_cols)
stride_x, # stride of each row in input
W_ptr, # pointer to weights, shape (n_cols,)
Mean_ptr, # pointer to mean, shape (n_rows,)
stride_mean, # stride of each row in mean
RSTD_ptr, # pointer to rstd, shape (n_rows,)
stride_rstd, # stride of each row in rstd
DX_ptr, # pointer to input grad, shape (n_rows, n_cols)
DW_ptr, # pointer to weights grad, shape (n_cols,)
DB_ptr, # pointer to bias grad, shape (n_cols,)
DY_ptr, # pointer to output grad, shape (n_rows, n_cols)
stride_x, # stride of each row in input
stride_dx, # stride of each row in input grad
DW_ptr, # pointer to weights grad, shape (n_cols,)
stride_dw, # stride of each row in weights grad
DB_ptr, # pointer to bias grad, shape (n_cols,)
stride_db, # stride of each row in bias grad
DY_ptr, # pointer to output grad, shape (n_rows, n_cols)
stride_dy, # stride of each row in output grad
n_rows,
n_cols,
rows_per_program: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
dtype: tl.constexpr,
atomic_dtype: tl.constexpr,
):
"""
References:
https://arxiv.org/abs/1607.06450
https://github.com/karpathy/llm.c/blob/master/doc/layernorm/layernorm.md
"""
row_idx = tl.program_id(0).to(tl.int64)
row_block_id = tl.program_id(0).to(tl.int64)
row_start = row_block_id * rows_per_program
row_end = min((row_block_id + 1) * rows_per_program, n_rows)
cols = tl.arange(0, BLOCK_SIZE)
mask = cols < n_cols
dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
db_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
# Pre-load weights once (same optimization as forward pass)
w = tl.load(W_ptr + cols, mask=mask, other=0.0)
w_f32 = w.to(tl.float32)
# Calculate pointers for this specific row
row_X_ptr = X_ptr + row_idx * stride_x
row_DX_ptr = DX_ptr + row_idx * stride_dx
row_DY_ptr = DY_ptr + row_idx * stride_dy
row_Mean_ptr = Mean_ptr + row_idx
row_RSTD_ptr = RSTD_ptr + row_idx
row_X_ptr = X_ptr + row_start * stride_x
row_DX_ptr = DX_ptr + row_start * stride_dx
row_DY_ptr = DY_ptr + row_start * stride_dy
row_Mean_ptr = Mean_ptr + row_start
row_RSTD_ptr = RSTD_ptr + row_start
# Load data for this row
x = tl.load(row_X_ptr + cols, mask=mask, other=0.0)
dy = tl.load(row_DY_ptr + cols, mask=mask, other=0.0)
mean = tl.load(row_Mean_ptr)
rstd = tl.load(row_RSTD_ptr)
for _ in range(row_start, row_end):
# Load data for this row
x = tl.load(row_X_ptr + cols, mask=mask, other=0.0)
dy = tl.load(row_DY_ptr + cols, mask=mask, other=0.0)
mean = tl.load(row_Mean_ptr)
rstd = tl.load(row_RSTD_ptr)
# Convert to fp32 for numerical stability
x_f32 = x.to(tl.float32)
dy_f32 = dy.to(tl.float32)
mean_f32 = mean.to(tl.float32)
rstd_f32 = rstd.to(tl.float32)
# Convert to fp32 for numerical stability
x_f32 = x.to(tl.float32)
dy_f32 = dy.to(tl.float32)
mean_f32 = mean.to(tl.float32)
rstd_f32 = rstd.to(tl.float32)
# Compute backward pass for this row
x_hat = (x_f32 - mean_f32) * rstd_f32
wdy = w_f32 * dy_f32
c1 = tl.sum(x_hat * wdy, axis=0) / n_cols
c2 = tl.sum(wdy, axis=0) / n_cols
dx = (wdy - (x_hat * c1 + c2)) * rstd_f32
# Compute backward pass for this row
x_hat = (x_f32 - mean_f32) * rstd_f32
wdy = w_f32 * dy_f32
c1 = tl.sum(x_hat * wdy, axis=0) / n_cols
c2 = tl.sum(wdy, axis=0) / n_cols
dx = (wdy - (x_hat * c1 + c2)) * rstd_f32
# Store input gradient
tl.store(row_DX_ptr + cols, dx.to(dtype), mask=mask)
# Store input gradient
tl.store(row_DX_ptr + cols, dx, mask=mask)
# Accumulate weight and bias gradients using atomic operations
dw = dy_f32 * x_hat
db = dy_f32
tl.atomic_add(DW_ptr + cols, dw.to(atomic_dtype), mask=mask)
tl.atomic_add(DB_ptr + cols, db.to(atomic_dtype), mask=mask)
# Accumulate weight and bias gradients for this thread block's assigned rows
dw = dy_f32 * x_hat
db = dy_f32
dW_row += dw
db_row += db
row_X_ptr += stride_x
row_DX_ptr += stride_dx
row_DY_ptr += stride_dy
row_Mean_ptr += stride_mean
row_RSTD_ptr += stride_rstd
tl.store(DW_ptr + row_block_id * stride_dw + cols, dW_row, mask=mask)
tl.store(DB_ptr + row_block_id * stride_db + cols, db_row, mask=mask)
def layer_norm_forward(X, W, B, eps):
@@ -228,31 +248,25 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
dY = dY.view(-1, dim)
n_rows, n_cols = dY.shape
# Allocate gradient tensors
DX = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
# Use float32 for weight/bias gradients if bfloat16 (due to atomic_add limitation)
grad_dtype = torch.float32 if W.dtype == torch.bfloat16 else W.dtype
DW = torch.zeros(n_cols, dtype=grad_dtype, device=W.device)
DB = torch.zeros(n_cols, dtype=grad_dtype, device=W.device)
sm_count = 1
if X.device.type == "cuda":
sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
elif X.device.type == "xpu":
sm_count = torch.xpu.get_device_properties(X.device).gpu_eu_count
# fp32 for numerical stability especially.
_DW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
_DB = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
# Calculate optimal block size and warp configuration
BLOCK_SIZE, num_warps = calculate_settings(n_cols)
if n_cols > BLOCK_SIZE:
raise RuntimeError(f"Feature dimension {n_cols} exceeds maximum supported size of {BLOCK_SIZE}.")
rows_per_program = math.ceil(n_rows / sm_count)
grid = (sm_count,)
# Determine dtype for triton operations
triton_dtype = (
tl.float32
if X.dtype == torch.float32
else tl.bfloat16
if X.dtype == torch.bfloat16
else tl.float16
if X.dtype == torch.float16
else tl.float32 # fallback
)
# Use float32 for atomic operations if bfloat16 is not supported
atomic_dtype = tl.float32 if triton_dtype == tl.bfloat16 else triton_dtype
# Allocate gradient tensors
DX = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
kernel_args = {"num_warps": num_warps}
# XPU-specific optimization
@@ -260,28 +274,33 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
kernel_args.update({"grf_mode": "large", "num_warps": 32, "num_stages": 4})
# Launch kernel with one thread block per row for optimal performance
grid = (n_rows,)
_layer_norm_backward_kernel[grid](
X,
X.stride(0),
W,
Mean,
Mean.stride(0),
RSTD,
RSTD.stride(0),
DX,
DW,
DB,
dY,
X.stride(0),
DX.stride(0),
_DW,
_DW.stride(0),
_DB,
_DB.stride(0),
dY,
dY.stride(0),
n_rows,
n_cols,
rows_per_program=rows_per_program,
BLOCK_SIZE=BLOCK_SIZE,
dtype=triton_dtype,
atomic_dtype=atomic_dtype,
**kernel_args,
)
DX = DX.view(*shape)
return DX, DW.to(W.dtype), DB.to(W.dtype)
DW = _DW.sum(dim=0).to(W.dtype)
DB = _DB.sum(dim=0).to(B.dtype)
return DX, DW, DB
class LigerLayerNormFunction(torch.autograd.Function):
+136
View File
@@ -0,0 +1,136 @@
import math
from typing import Callable
from typing import List
from typing import Optional
import torch
from liger_kernel.ops.utils import ensure_contiguous
class LigerTiledMLPFunction(torch.autograd.Function):
"""
Based on DeepSpeed's TiledMLP:
https://github.com/deepspeedai/DeepSpeed/blob/v0.18.2/deepspeed/runtime/sequence_parallel/ulysses_sp.py#L838
Perform a tiled MLP computation to massively reduce memory usage needed to compute MLP
when using very long sequence lengths.
This module re-computes `forward` in the `backward`. So the `forward` occurs twice each iteration.
And if you're using activation checkpointing it then occurs thrice.
Args:
fn: the function to call on sharded inputs (e.g., mlp.forward)
mlp_module: the MLP nn.Module object
x: the input to MLP.forward (hidden_states)
shards: how many shards to use
compute_params: a list of weights engaged in the compute
Returns:
the computed hidden_states
"""
@staticmethod
@ensure_contiguous
def forward(
ctx,
fn: Callable,
mlp_module: torch.nn.Module,
x: torch.Tensor,
shards: int,
compute_params: Optional[List[torch.nn.Parameter]] = None,
) -> torch.Tensor:
ctx.fn = fn
ctx.mlp_module = mlp_module
ctx.shards = shards
ctx.save_for_backward(x)
# x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size] (moe experts)
x_shards = list(torch.chunk(x, chunks=shards, dim=-2))
with torch.no_grad():
output_shards = [fn(mlp_module, x_shard) for x_shard in x_shards]
output_unsharded = torch.cat(output_shards, dim=-2)
return output_unsharded
@staticmethod
@ensure_contiguous
def backward(ctx, *grads) -> tuple:
fn = ctx.fn
(x,) = ctx.saved_tensors
mlp_module = ctx.mlp_module
shards = ctx.shards
x_requires_grad = x.requires_grad
x = x.detach()
# detach() unsets x.requires_grad, so restore it
x.requires_grad_(x_requires_grad)
# x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size] (moe experts)
hidden_size = x.shape[-1]
x_shape_orig = x.shape
# flatten bs+seqlen to avoid having stride issues when narrowing into seqlen w/ bs>1
x = x.view(-1, hidden_size)
incoming_grad = grads[0].view(-1, hidden_size)
x_grad = torch.zeros_like(x)
x_shards = list(torch.chunk(x, chunks=shards, dim=0))
for i, x_shard in enumerate(x_shards):
x_shard.requires_grad_(x_requires_grad)
# if seqlen is not exactly divisible by shards the last step will be shorter than shard_step
shard_step = x_shards[i].shape[0]
shard_offset = i * x_shards[0].shape[0]
x_shard.grad = x_grad.narrow(0, shard_offset, shard_step).view_as(x_shard)
incoming_grad_shard = incoming_grad.narrow(0, shard_offset, shard_step).view_as(x_shard)
with torch.enable_grad():
output = fn(mlp_module, x_shard)
torch.autograd.backward(output, incoming_grad_shard)
# unflatten
x_grad = x_grad.view(x_shape_orig)
return (None, None, x_grad, None, None)
def apply_tiled_mlp(
fn: Callable,
mlp_module: torch.nn.Module,
x: torch.Tensor,
num_shards: Optional[int] = None,
compute_params: Optional[List[torch.nn.Parameter]] = None,
) -> torch.Tensor:
"""
Apply tiled MLP computation for memory efficiency.
Args:
fn: the function to call on sharded inputs (e.g., lambda module, x: module(x))
mlp_module: the MLP nn.Module object
x: the input tensor with shape [bs, seqlen, hidden_size] or [seqlen, hidden_size]
num_shards: number of shards to use. If None, automatically calculated as ceil(seqlen / hidden_size)
compute_params: list of parameters for DeepSpeed ZeRO optimization
Returns:
output tensor with the same shape as input
"""
if num_shards is None:
# x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size]
hidden_size = x.shape[-1]
seqlen = x.shape[-2]
num_shards = math.ceil(seqlen / hidden_size)
# Ensure num_shards is at least 1
num_shards = max(1, num_shards)
return LigerTiledMLPFunction.apply(
fn,
mlp_module,
x,
num_shards,
compute_params,
)
+13
View File
@@ -24,6 +24,8 @@ from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP # noqa: F4
from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP # noqa: F401
from liger_kernel.transformers.swiglu import LigerQwen3MoeSwiGLUMLP # noqa: F401
from liger_kernel.transformers.swiglu import LigerSwiGLUMLP # noqa: F401
from liger_kernel.transformers.tiled_mlp import LigerTiledGEGLUMLP # noqa: F401
from liger_kernel.transformers.tiled_mlp import LigerTiledSwiGLUMLP # noqa: F401
from liger_kernel.transformers.tvd import LigerTVDLoss # noqa: F401
# Static-only imports for IDEs and type checkers
@@ -40,6 +42,8 @@ if TYPE_CHECKING:
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4v # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4v_moe # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_granite # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_hunyuan_v1_dense # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_hunyuan_v1_moe # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_internvl # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama4 # noqa: F401
@@ -48,6 +52,7 @@ if TYPE_CHECKING:
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mixtral # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mllama # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_olmo2 # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_olmo3 # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_paligemma # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_phi3 # noqa: F401
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_qwen2 # noqa: F401
@@ -114,6 +119,7 @@ def __getattr__(name: str):
"apply_liger_kernel_to_mixtral",
"apply_liger_kernel_to_mllama",
"apply_liger_kernel_to_olmo2",
"apply_liger_kernel_to_olmo3",
"apply_liger_kernel_to_paligemma",
"apply_liger_kernel_to_phi3",
"apply_liger_kernel_to_qwen2",
@@ -126,6 +132,8 @@ def __getattr__(name: str):
"apply_liger_kernel_to_qwen3_vl_moe",
"apply_liger_kernel_to_smollm3",
"apply_liger_kernel_to_smolvlm",
"apply_liger_kernel_to_hunyuan_v1_dense",
"apply_liger_kernel_to_hunyuan_v1_moe",
}
if name in monkey_patch_symbols:
@@ -155,6 +163,8 @@ __all__ = [
"LigerPhi3SwiGLUMLP",
"LigerQwen3MoeSwiGLUMLP",
"LigerSwiGLUMLP",
"LigerTiledGEGLUMLP",
"LigerTiledSwiGLUMLP",
"LigerTVDLoss",
"LigerKLDIVLoss",
"LigerMultiTokenAttention",
@@ -186,6 +196,7 @@ if _TRANSFORMERS_AVAILABLE:
"apply_liger_kernel_to_mixtral",
"apply_liger_kernel_to_mllama",
"apply_liger_kernel_to_olmo2",
"apply_liger_kernel_to_olmo3",
"apply_liger_kernel_to_paligemma",
"apply_liger_kernel_to_phi3",
"apply_liger_kernel_to_qwen2",
@@ -198,5 +209,7 @@ if _TRANSFORMERS_AVAILABLE:
"apply_liger_kernel_to_qwen3_vl_moe",
"apply_liger_kernel_to_smollm3",
"apply_liger_kernel_to_smolvlm",
"apply_liger_kernel_to_hunyuan_v1_dense",
"apply_liger_kernel_to_hunyuan_v1_moe",
]
)
+56 -1
View File
@@ -1,3 +1,6 @@
import torch
from liger_kernel.chunked_loss.fused_linear_ppo import LigerFusedLinearPPOBase
from liger_kernel.ops.grpo_loss import GrpoLossFunction
@@ -13,12 +16,20 @@ def triton_grpo_loss(
eps_low=0.2,
eps_high=0.4,
inplace=True,
loss_type="dapo",
max_completion_length=None,
importance_sampling_level="token",
reduce=False,
):
assert logits is not None and completion_ids is not None and advantages is not None, (
"must provide logits、completion_ids and advantages"
)
if importance_sampling_level != "token":
raise ValueError(
f"Triton GRPO loss only supports token-level importance sampling. Got {importance_sampling_level}."
)
return GrpoLossFunction.apply(
per_token_loss, per_token_kl, is_clipped = GrpoLossFunction.apply(
logits,
old_logp,
ref_logp,
@@ -31,6 +42,50 @@ def triton_grpo_loss(
eps_high,
inplace,
)
if not reduce:
return per_token_loss, per_token_kl, is_clipped
loss = _reduce_grpo_loss(
per_token_loss,
completion_mask,
loss_type=loss_type,
max_completion_length=max_completion_length,
)
metrics = []
if beta != 0.0 and per_token_kl is not None:
metrics.append(_masked_mean(per_token_kl, completion_mask))
metrics.append(_masked_mean(is_clipped.float(), completion_mask))
return loss, metrics
def _reduce_grpo_loss(per_token_loss, completion_mask, loss_type, max_completion_length):
mask = completion_mask
if mask is None:
mask = torch.ones_like(per_token_loss, dtype=per_token_loss.dtype, device=per_token_loss.device)
mask = mask.to(per_token_loss.dtype)
if loss_type == "grpo":
per_seq = (per_token_loss * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)
return per_seq.mean()
if loss_type == "bnpo":
return (per_token_loss * mask).sum() / mask.sum().clamp(min=1.0)
if loss_type == "dr_grpo":
if max_completion_length is None:
raise ValueError("max_completion_length must be provided when using loss_type='dr_grpo'")
batch = per_token_loss.shape[0]
return (per_token_loss * mask).sum() / (batch * max_completion_length)
if loss_type == "dapo":
normalizer = LigerFusedLinearPPOBase._compute_dapo_normalizer(mask)
return (per_token_loss * mask).sum() / normalizer
raise ValueError(f"Unsupported loss_type '{loss_type}' for Triton GRPO loss.")
def _masked_mean(values, mask):
if mask is None:
mask = torch.ones_like(values, dtype=values.dtype, device=values.device)
mask = mask.to(values.dtype)
return (values * mask).sum() / mask.sum().clamp(min=1.0)
# This is a demo how to use grpo_loss in GRPOTrainer. The Trl version must be 0.16
@@ -0,0 +1,134 @@
from typing import List
from typing import Optional
from typing import Union
import torch
from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
from liger_kernel.transformers.model.loss_utils import unpack_cross_entropy_result
from liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWithPast
def lce_forward(
self,
input_ids: Optional[torch.LongTensor] = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[List[torch.FloatTensor]] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
labels: Optional[torch.LongTensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
cache_position: Optional[torch.LongTensor] = None,
logits_to_keep: Union[int, torch.Tensor] = 0,
skip_logits: Optional[bool] = None,
return_dict: Optional[bool] = None,
**kwargs,
) -> LigerCausalLMOutputWithPast:
r"""
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
logits_to_keep (`int` or `torch.Tensor`, *optional*):
If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
`input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
This is useful when using packed tensor format (single dimension for batch and sequence length).
Returns:
Example:
```python
>>> from transformers import AutoTokenizer, HunYuanDenseV1ForCausalLM
>>> model = HunYuanDenseV1ForCausalLM.from_pretrained("meta-hunyuan_v1_dense/HunYuanDenseV1-2-7b-hf")
>>> tokenizer = AutoTokenizer.from_pretrained("meta-hunyuan_v1_dense/HunYuanDenseV1-2-7b-hf")
>>> prompt = "Hey, are you conscious? Can you talk to me?"
>>> inputs = tokenizer(prompt, return_tensors="pt")
>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
```"""
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
outputs = self.model(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
cache_position=cache_position,
**kwargs,
)
hidden_states = outputs[0]
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
kept_hidden_states = hidden_states[:, slice_indices, :]
shift_labels = kwargs.pop("shift_labels", None)
logits = None
loss = None
token_accuracy = None
if skip_logits and labels is None and shift_labels is None:
raise ValueError("skip_logits is True, but labels and shift_labels are None")
if skip_logits is None:
# By default, if in training mode, don't materialize logits
skip_logits = self.training and (labels is not None or shift_labels is not None)
# Compute loss
if skip_logits:
result = LigerForCausalLMLoss(
hidden_states=kept_hidden_states,
lm_head_weight=self.lm_head.weight,
labels=labels,
shift_labels=shift_labels,
hidden_size=self.config.hidden_size,
**kwargs,
)
loss, _, token_accuracy = unpack_cross_entropy_result(result)
else:
logits = self.lm_head(kept_hidden_states)
if labels is not None or shift_labels is not None:
loss = self.loss_function(
logits=logits,
labels=labels,
shift_labels=shift_labels,
vocab_size=self.config.vocab_size,
**kwargs,
)
if not return_dict:
output = (logits,) + outputs[1:]
output = ((loss,) + output) if loss is not None else output
output = output + (token_accuracy,) if token_accuracy is not None else output
return output
# Return custom output class with accuracy field
return LigerCausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
token_accuracy=token_accuracy,
)
@@ -0,0 +1,142 @@
from typing import List
from typing import Optional
from typing import Tuple
from typing import Union
import torch
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.utils.deprecation import deprecate_kwarg
from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
from liger_kernel.transformers.model.loss_utils import unpack_cross_entropy_result
from liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWithPast
@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
def lce_forward(
self,
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[List[torch.FloatTensor]] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
labels: Optional[torch.LongTensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
cache_position: Optional[torch.LongTensor] = None,
logits_to_keep: Union[int, torch.Tensor] = 0,
skip_logits: Optional[bool] = None,
**kwargs,
) -> Union[Tuple, LigerCausalLMOutputWithPast]:
r"""
Args:
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
logits_to_keep (`int` or `torch.Tensor`, *optional*):
If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
`input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
This is useful when using packed tensor format (single dimension for batch and sequence length).
Returns:
Example:
```python
>>> from transformers import AutoTokenizer, Olmo3ForCausalLM
>>> model = Olmo3ForCausalLM.from_pretrained("allenai/Olmo-3-7B-Instruct")
>>> tokenizer = AutoTokenizer.from_pretrained("allenai/Olmo-3-7B-Instruct")
>>> prompt = "Hey, are you conscious? Can you talk to me?"
>>> inputs = tokenizer(prompt, return_tensors="pt")
>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
'Hey, are you conscious? Can you talk to me?\nI’m not sure if you’re conscious of this, but I’m'
```
"""
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
outputs: BaseModelOutputWithPast = self.model(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
cache_position=cache_position,
**kwargs,
)
hidden_states = outputs[0]
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
kept_hidden_states = hidden_states[:, slice_indices, :]
shift_labels = kwargs.pop("shift_labels", None)
logits = None
loss = None
token_accuracy = None
if skip_logits and labels is None and shift_labels is None:
raise ValueError("skip_logits is True, but labels and shift_labels are None")
if skip_logits is None:
# By default, if in training mode, don't materialize logits
skip_logits = self.training and (labels is not None or shift_labels is not None)
# Compute loss
if skip_logits:
result = LigerForCausalLMLoss(
hidden_states=kept_hidden_states,
lm_head_weight=self.lm_head.weight,
labels=labels,
shift_labels=shift_labels,
hidden_size=self.config.hidden_size,
**kwargs,
)
loss, _, token_accuracy = unpack_cross_entropy_result(result)
else:
logits = self.lm_head(kept_hidden_states)
if labels is not None or shift_labels is not None:
loss = self.loss_function(
logits=logits,
labels=labels,
shift_labels=shift_labels,
vocab_size=self.config.vocab_size,
**kwargs,
)
if not return_dict:
output = (logits,) + outputs[1:]
output = ((loss,) + output) if loss is not None else output
output = output + (token_accuracy,) if token_accuracy is not None else output
return output
# Return custom output class with token_accuracy field
return LigerCausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
token_accuracy=token_accuracy,
)
@@ -83,6 +83,8 @@ def lce_forward(
kept_hidden_states = hidden_states[:, slice_indices, :]
shift_labels = kwargs.pop("shift_labels", None)
# Remove output-control parameters that shouldn't be passed to loss functions
kwargs.pop("return_dict", None)
logits = None
loss = None
token_accuracy = None
@@ -1928,6 +1928,74 @@ def apply_liger_kernel_to_olmo2(
_patch_rms_norm_module(decoder_layer.post_feedforward_layernorm, in_place=False)
def apply_liger_kernel_to_olmo3(
rope: bool = True,
cross_entropy: bool = False,
fused_linear_cross_entropy: bool = True,
rms_norm: bool = True,
swiglu: bool = True,
model: PreTrainedModel = None,
) -> None:
"""
Apply Liger kernels to replace original implementation in HuggingFace Olmo3 models.
Args:
rope (bool): Whether to apply Liger's rotary position embedding. Default is True.
cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False.
fused_linear_cross_entropy (bool):
Whether to apply Liger's fused linear cross entropy loss. Default is True.
`cross_entropy` and `fused_linear_cross_entropy` cannot both be True.
If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
rms_norm (bool): Whether to apply Liger's RMSNorm. Default is True.
swiglu (bool): Whether to apply Liger's SwiGLU to Olmo3MLP. Default is True.
model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
loaded. Default is None.
"""
assert not (cross_entropy and fused_linear_cross_entropy), (
"cross_entropy and fused_linear_cross_entropy cannot both be True."
)
from transformers.models.olmo3 import modeling_olmo3
from transformers.models.olmo3.modeling_olmo3 import Olmo3Model
from liger_kernel.transformers.model.olmo3 import lce_forward as olmo3_lce_forward
from liger_kernel.transformers.rms_norm import LigerRMSNormForOlmo2
# Olmo3 arch is very similar to Olmo2, so we can reuse all these components in the same way.
if rope:
modeling_olmo3.apply_rotary_pos_emb = liger_rotary_pos_emb
if rms_norm:
modeling_olmo3.Olmo3RMSNorm = LigerRMSNormForOlmo2 # same as olmo2
if swiglu:
modeling_olmo3.Olmo3MLP = LigerSwiGLUMLP
if cross_entropy:
from transformers.loss.loss_utils import nn
nn.functional.cross_entropy = liger_cross_entropy
if fused_linear_cross_entropy:
if model is not None:
model.forward = MethodType(olmo3_lce_forward, model)
else:
modeling_olmo3.Olmo3ForCausalLM.forward = olmo3_lce_forward
if model is not None:
# The model instance already exists, so we need to additionally patch the
# instance variables that reference already-instantiated modules
# get the base model from the model instance
base_model: Olmo3Model = getattr(model, model.base_model_prefix, model)
if rms_norm:
_patch_rms_norm_module(base_model.norm)
for decoder_layer in base_model.layers:
if swiglu:
_patch_swiglu_module(decoder_layer.mlp, LigerSwiGLUMLP)
if rms_norm:
_patch_rms_norm_module(decoder_layer.post_attention_layernorm, in_place=False)
_patch_rms_norm_module(decoder_layer.post_feedforward_layernorm, in_place=False)
def apply_liger_kernel_to_glm4(
rope: bool = False,
cross_entropy: bool = False,
@@ -2558,6 +2626,123 @@ def apply_liger_kernel_to_qwen3_next(
_patch_swiglu_module(expert, LigerQwen3MoeSwiGLUMLP)
def apply_liger_kernel_to_hunyuan_v1_dense(
rope: bool = True,
cross_entropy: bool = False,
fused_linear_cross_entropy: bool = True,
rms_norm: bool = True,
swiglu: bool = True,
model: PreTrainedModel = None,
) -> None:
"""
Apply Liger kernels to replace original implementation in HuggingFace Hunyuan v1 dense models.
"""
assert not (cross_entropy and fused_linear_cross_entropy), (
"cross_entropy and fused_linear_cross_entropy cannot both be True."
)
from transformers.models.hunyuan_v1_dense import modeling_hunyuan_v1_dense
from transformers.models.hunyuan_v1_dense.modeling_hunyuan_v1_dense import HunYuanDenseV1Model
from liger_kernel.transformers.model.hunyuan_v1 import lce_forward as hunyuan_v1_lce_forward
from liger_kernel.transformers.swiglu import LigerHunyuanV1SwiGLUMLP
if rope:
modeling_hunyuan_v1_dense.apply_rotary_pos_emb = liger_rotary_pos_emb
if rms_norm:
modeling_hunyuan_v1_dense.HunYuanDenseV1RMSNorm = LigerRMSNorm
if cross_entropy:
from transformers.loss.loss_utils import nn
nn.functional.cross_entropy = liger_cross_entropy
if fused_linear_cross_entropy:
if model is not None:
model.forward = MethodType(hunyuan_v1_lce_forward, model)
else:
modeling_hunyuan_v1_dense.HunYuanDenseV1ForCausalLM.forward = hunyuan_v1_lce_forward
if swiglu:
modeling_hunyuan_v1_dense.HunYuanDenseV1MLP = LigerHunyuanV1SwiGLUMLP
if model is not None:
# The model instance already exists, so we need to additionally patch the
# instance variables that reference already-instantiated modules
# get the base model from the model instance
base_model: HunYuanDenseV1Model = getattr(model, model.base_model_prefix, model)
if rms_norm:
_patch_rms_norm_module(base_model.norm)
for decoder_layer in base_model.layers:
if swiglu:
_patch_swiglu_module(decoder_layer.mlp, LigerHunyuanV1SwiGLUMLP)
if rms_norm:
_patch_rms_norm_module(decoder_layer.input_layernorm)
_patch_rms_norm_module(decoder_layer.post_attention_layernorm)
def apply_liger_kernel_to_hunyuan_v1_moe(
rope: bool = True,
cross_entropy: bool = False,
fused_linear_cross_entropy: bool = True,
rms_norm: bool = True,
swiglu: bool = True,
model: PreTrainedModel = None,
) -> None:
"""
Apply Liger kernels to replace original implementation in HuggingFace Qwen3 models.
"""
assert not (cross_entropy and fused_linear_cross_entropy), (
"cross_entropy and fused_linear_cross_entropy cannot both be True."
)
from transformers.models.hunyuan_v1_moe import modeling_hunyuan_v1_moe
from transformers.models.hunyuan_v1_moe.modeling_hunyuan_v1_moe import HunYuanMoEV1Model
from liger_kernel.transformers.model.hunyuan_v1 import lce_forward as hunyuan_v1_moe_lce_forward
from liger_kernel.transformers.swiglu import LigerHunyuanV1SwiGLUMLP
if rope:
modeling_hunyuan_v1_moe.apply_rotary_pos_emb = liger_rotary_pos_emb
if rms_norm:
modeling_hunyuan_v1_moe.HunYuanMoEV1RMSNorm = LigerRMSNorm
if cross_entropy:
from transformers.loss.loss_utils import nn
nn.functional.cross_entropy = liger_cross_entropy
if fused_linear_cross_entropy:
if model is not None:
model.forward = MethodType(hunyuan_v1_moe_lce_forward, model)
else:
modeling_hunyuan_v1_moe.HunYuanMoEV1ForCausalLM.forward = hunyuan_v1_moe_lce_forward
if swiglu:
modeling_hunyuan_v1_moe.HunYuanMoEV1MLP = LigerHunyuanV1SwiGLUMLP
if model is not None:
# The model instance already exists, so we need to additionally patch the
# instance variables that reference already-instantiated modules
# get the base model from the model instance
base_model: HunYuanMoEV1Model = getattr(model, model.base_model_prefix, model)
if rms_norm:
_patch_rms_norm_module(base_model.norm)
for decoder_layer in base_model.layers:
if swiglu:
for mlp_expert in decoder_layer.mlp.experts:
_patch_swiglu_module(mlp_expert, LigerHunyuanV1SwiGLUMLP)
if rms_norm:
_patch_rms_norm_module(decoder_layer.input_layernorm)
_patch_rms_norm_module(decoder_layer.post_attention_layernorm)
# Model type corresponds to the keys defined in transformers/models/auto/modeling_auto.py
MODEL_TYPE_TO_APPLY_LIGER_FN = {
"gemma": apply_liger_kernel_to_gemma,
@@ -2578,6 +2763,7 @@ MODEL_TYPE_TO_APPLY_LIGER_FN = {
"mistral": apply_liger_kernel_to_mistral,
"mixtral": apply_liger_kernel_to_mixtral,
"olmo2": apply_liger_kernel_to_olmo2,
"olmo3": apply_liger_kernel_to_olmo3,
"qwen2": apply_liger_kernel_to_qwen2,
"qwen3": apply_liger_kernel_to_qwen3,
"qwen3_moe": apply_liger_kernel_to_qwen3_moe,
@@ -2595,6 +2781,8 @@ MODEL_TYPE_TO_APPLY_LIGER_FN = {
"paligemma": apply_liger_kernel_to_paligemma,
"falcon_h1": apply_liger_kernel_to_falcon_h1,
"smolvlm": apply_liger_kernel_to_smolvlm,
"hunyuan_v1_dense": apply_liger_kernel_to_hunyuan_v1_dense,
"hunyuan_v1_moe": apply_liger_kernel_to_hunyuan_v1_moe,
}
+17
View File
@@ -77,3 +77,20 @@ class LigerQwen3MoeSwiGLUMLP(nn.Module):
def forward(self, x):
return self.down_proj(LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x)))
class LigerHunyuanV1SwiGLUMLP(nn.Module):
def __init__(self, config, layer_idx=None, is_shared_mlp=False):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size
self.intermediate_size = config.intermediate_size
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
self.layer_idx = layer_idx
if config.hidden_act not in ["silu", "swish"]:
raise ValueError(f"Activation function {config.hidden_act} not supported.")
def forward(self, x):
return self.down_proj(LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x)))
+133
View File
@@ -0,0 +1,133 @@
from typing import Optional
import torch.nn as nn
from liger_kernel.ops.geglu import LigerGELUMulFunction
from liger_kernel.ops.swiglu import LigerSiLUMulFunction
from liger_kernel.ops.tiled_mlp import apply_tiled_mlp
class LigerTiledGEGLUMLP(nn.Module):
"""
Memory-efficient GEGLU MLP using tiled computation.
This module combines GEGLU activation with tiled processing to handle
very long sequences efficiently. The forward pass is recomputed during
backward to save memory.
Args:
config: Model configuration with hidden_size and intermediate_size attributes
num_shards: Number of shards to split the sequence. If None, automatically
calculated as ceil(seqlen / hidden_size)
"""
def __init__(self, config, num_shards: Optional[int] = None):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size
self.intermediate_size = config.intermediate_size
self.num_shards = num_shards
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
# Validate activation function
if hasattr(config, "hidden_act") and config.hidden_act not in [
"gelu",
"gelu_new",
"gelu_pytorch_tanh",
]:
raise ValueError(f"LigerTiledGEGLUMLP requires GELU activation, got {config.hidden_act}")
def _mlp_forward(self, module, x):
"""Internal MLP forward function for tiled computation."""
gate = module.gate_proj(x)
up = module.up_proj(x)
return module.down_proj(LigerGELUMulFunction.apply(gate, up))
def forward(self, x):
"""
Forward pass with tiled computation.
Args:
x: Input tensor of shape [batch_size, seq_len, hidden_size]
or [seq_len, hidden_size]
Returns:
Output tensor of the same shape as input
"""
compute_params = [
self.gate_proj.weight,
self.up_proj.weight,
self.down_proj.weight,
]
return apply_tiled_mlp(
fn=self._mlp_forward,
mlp_module=self,
x=x,
num_shards=self.num_shards,
compute_params=compute_params,
)
class LigerTiledSwiGLUMLP(nn.Module):
"""
Memory-efficient SwiGLU MLP using tiled computation.
This module combines SwiGLU activation with tiled processing to handle
very long sequences efficiently. The forward pass is recomputed during
backward to save memory.
Args:
config: Model configuration with hidden_size and intermediate_size attributes
num_shards: Number of shards to split the sequence. If None, automatically
calculated as ceil(seqlen / hidden_size)
"""
def __init__(self, config, num_shards: Optional[int] = None):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size
self.intermediate_size = config.intermediate_size
self.num_shards = num_shards
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
# Validate activation function
if hasattr(config, "hidden_act") and config.hidden_act not in ["silu", "swish"]:
raise ValueError(f"LigerTiledSwiGLUMLP requires SiLU/Swish activation, got {config.hidden_act}")
def _mlp_forward(self, module, x):
"""Internal MLP forward function for tiled computation."""
gate = module.gate_proj(x)
up = module.up_proj(x)
return module.down_proj(LigerSiLUMulFunction.apply(gate, up))
def forward(self, x):
"""
Forward pass with tiled computation.
Args:
x: Input tensor of shape [batch_size, seq_len, hidden_size]
or [seq_len, hidden_size]
Returns:
Output tensor of the same shape as input
"""
compute_params = [
self.gate_proj.weight,
self.up_proj.weight,
self.down_proj.weight,
]
return apply_tiled_mlp(
fn=self._mlp_forward,
mlp_module=self,
x=x,
num_shards=self.num_shards,
compute_params=compute_params,
)
+224 -47
View File
@@ -5,6 +5,8 @@ import torch.nn.functional as F
from liger_kernel.chunked_loss import LigerFusedLinearGRPOLoss
from liger_kernel.chunked_loss.functional import liger_fused_linear_grpo
from liger_kernel.chunked_loss.grpo_loss import LigerFusedLinearGRPOFunction
from liger_kernel.transformers.grpo_loss import _reduce_grpo_loss
from liger_kernel.transformers.grpo_loss import triton_grpo_loss
from liger_kernel.utils import infer_device
from test.utils import assert_verbose_allclose
from test.utils import set_seed
@@ -45,6 +47,60 @@ class TorchLMHeadGRPO(torch.nn.Module):
if self.loss_type == "dr_grpo":
assert self.max_completion_length is not None, "max_completion_length must be provided for dr_grpo"
@staticmethod
def compute_per_token_components(
per_token_logps,
attention_mask,
advantages,
old_per_token_logps,
ref_per_token_logps,
epsilon_low,
epsilon_high,
beta,
importance_sampling_level,
):
attention_mask = attention_mask.to(per_token_logps.dtype)
old_per_token_logps = (
old_per_token_logps.float() if old_per_token_logps is not None else per_token_logps.detach()
)
log_ratio = per_token_logps - old_per_token_logps
if importance_sampling_level == "token":
log_importance_weights = log_ratio
elif importance_sampling_level == "sequence":
log_importance_weights = (log_ratio * attention_mask).sum(-1) / attention_mask.sum(-1).clamp(min=1.0)
log_importance_weights = log_importance_weights.unsqueeze(-1)
else:
raise ValueError(
f"Unknown importance sampling level: {importance_sampling_level}. Possible values are 'token' "
"and 'sequence'."
)
coef_1 = torch.exp(log_importance_weights)
coef_2 = torch.clamp(coef_1, 1 - epsilon_low, 1 + epsilon_high)
expanded_advantages = advantages.unsqueeze(1)
per_token_loss1 = coef_1 * expanded_advantages
per_token_loss2 = coef_2 * expanded_advantages
per_token_loss = -torch.min(per_token_loss1, per_token_loss2)
kl_div = None
if beta != 0.0:
ref_per_token_logps = ref_per_token_logps.float()
kl_div = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1.0
per_token_loss = per_token_loss + beta * kl_div
if importance_sampling_level == "token":
is_clipped = ((coef_1 < 1 - epsilon_low) & (expanded_advantages < 0)) | (
(coef_1 > 1 + epsilon_high) & (expanded_advantages > 0)
)
else: # sequence level
# For sequence level, coef_1 is shape (B, 1), advantages is shape (B,)
seq_advantages = advantages
is_clipped = ((coef_1.squeeze(-1) < 1 - epsilon_low) & (seq_advantages < 0)) | (
(coef_1.squeeze(-1) > 1 + epsilon_high) & (seq_advantages > 0)
)
is_clipped = is_clipped.unsqueeze(1).expand_as(attention_mask)
return per_token_loss, kl_div, is_clipped
def forward(
self,
x, # Shape: [batch_size, seq_len, hidden_size]
@@ -82,33 +138,17 @@ class TorchLMHeadGRPO(torch.nn.Module):
else:
ref_per_token_logps = per_token_logps.detach()
# Compute policy gradient loss with importance sampling ratio
old_per_token_logps = (
old_per_token_logps.float() if old_per_token_logps is not None else per_token_logps.detach()
per_token_loss, kl_div, is_clipped = self.compute_per_token_components(
per_token_logps,
attention_mask,
advantages,
old_per_token_logps,
ref_per_token_logps,
self.epsilon_low,
self.epsilon_high,
self.beta,
self.importance_sampling_level,
)
log_ratio = per_token_logps - old_per_token_logps
if self.importance_sampling_level == "token":
log_importance_weights = log_ratio
elif self.importance_sampling_level == "sequence":
log_importance_weights = (log_ratio * attention_mask).sum(-1) / attention_mask.sum(-1).clamp(min=1.0)
log_importance_weights = log_importance_weights.unsqueeze(-1)
else:
raise ValueError(
f"Unknown importance sampling level: {self.importance_sampling_level}. Possible values are 'token' "
"and 'sequence'."
)
coef_1 = torch.exp(log_importance_weights)
coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high)
per_token_loss1 = coef_1 * advantages.unsqueeze(1)
per_token_loss2 = coef_2 * advantages.unsqueeze(1)
per_token_loss = -torch.min(per_token_loss1, per_token_loss2)
if self.beta != 0.0:
# Compute KL divergence between model and reference model
ref_per_token_logps = ref_per_token_logps.float()
kl_div = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1.0
per_token_loss = per_token_loss + self.beta * kl_div
# Apply masking and calculate loss based on loss_type
if self.loss_type == "grpo":
@@ -117,6 +157,9 @@ class TorchLMHeadGRPO(torch.nn.Module):
loss = (per_token_loss * attention_mask).sum() / torch.clamp(attention_mask.sum(), min=1.0)
elif self.loss_type == "dr_grpo":
loss = (per_token_loss * attention_mask).sum() / (per_token_loss.size(0) * self.max_completion_length)
elif self.loss_type == "dapo":
normalizer = attention_mask.sum().clamp(min=1.0)
loss = (per_token_loss * attention_mask).sum() / normalizer
else:
raise ValueError(f"Unknown loss type: {self.loss_type}")
@@ -124,18 +167,7 @@ class TorchLMHeadGRPO(torch.nn.Module):
metrics = []
if self.beta != 0.0:
metrics.append(((kl_div * attention_mask).sum() / torch.clamp(attention_mask.sum(), min=1.0)))
# Adjust clipping metric calculation based on importance sampling level
if self.importance_sampling_level == "token":
is_clipped = ((coef_1 < 1 - self.epsilon_low) & (advantages.unsqueeze(1) < 0)) | (
(coef_1 > 1 + self.epsilon_high) & (advantages.unsqueeze(1) > 0)
)
else: # sequence level
# For sequence level, coef_1 is shape (B, 1), advantages is shape (B,)
is_clipped = ((coef_1.squeeze(-1) < 1 - self.epsilon_low) & (advantages < 0)) | (
(coef_1.squeeze(-1) > 1 + self.epsilon_high) & (advantages > 0)
)
is_clipped = is_clipped.unsqueeze(1).expand_as(attention_mask)
metrics.append((is_clipped * attention_mask).sum() / torch.clamp(attention_mask.sum(), min=1.0))
metrics.append((is_clipped.float() * attention_mask).sum() / torch.clamp(attention_mask.sum(), min=1.0))
return loss, metrics
@@ -227,7 +259,7 @@ class LigerLMHeadGRPO(torch.nn.Module):
(False, False, True),
],
)
@pytest.mark.parametrize("loss_type", ["bnpo", "grpo", "dr_grpo"])
@pytest.mark.parametrize("loss_type", ["bnpo", "grpo", "dr_grpo", "dapo"])
@pytest.mark.parametrize("importance_sampling_level", ["token", "sequence"])
def test_correctness(
B,
@@ -515,12 +547,157 @@ def test_functional_correctness(
for metric1, metric2 in zip(aux1, aux2):
assert_verbose_allclose(metric1, metric2, atol=atol, rtol=rtol)
# Backward pass
loss1.backward()
loss2.backward()
# Check gradients match
assert_verbose_allclose(input1.grad, input2.grad, atol=atol, rtol=rtol)
assert_verbose_allclose(weight1.grad, weight2.grad, atol=atol, rtol=rtol)
if bias:
assert_verbose_allclose(bias1.grad, bias2.grad, atol=atol, rtol=rtol)
@pytest.mark.parametrize("loss_type", ["grpo", "bnpo", "dr_grpo", "dapo"])
def test_reduce_grpo_loss_matches_reference(loss_type):
torch.manual_seed(0)
per_token_loss = torch.randn(3, 5)
mask = torch.randint(0, 2, (3, 5), device=per_token_loss.device, dtype=torch.long)
mask[:, 0] = 1 # ensure at least one valid token per sequence
max_completion_length = 5 if loss_type == "dr_grpo" else None
reduced = _reduce_grpo_loss(per_token_loss, mask, loss_type, max_completion_length)
mask_f = mask.to(per_token_loss.dtype)
if loss_type == "grpo":
expected = ((per_token_loss * mask_f).sum(-1) / mask_f.sum(-1).clamp(min=1.0)).mean()
elif loss_type == "bnpo":
expected = (per_token_loss * mask_f).sum() / mask_f.sum().clamp(min=1.0)
elif loss_type == "dr_grpo":
expected = (per_token_loss * mask_f).sum() / (per_token_loss.size(0) * max_completion_length)
else: # dapo
expected = (per_token_loss * mask_f).sum() / mask_f.sum().clamp(min=1.0)
assert_verbose_allclose(reduced, expected)
def test_reduce_grpo_loss_requires_max_completion_length():
per_token_loss = torch.randn(2, 3)
mask = torch.ones_like(per_token_loss, dtype=torch.long)
with pytest.raises(ValueError):
_reduce_grpo_loss(per_token_loss, mask, "dr_grpo", max_completion_length=None)
@pytest.mark.parametrize("loss_type,beta", [("bnpo", 0.0), ("dapo", 0.04)])
def test_triton_grpo_loss_matches_reference(loss_type, beta):
pytest.importorskip("triton")
device = infer_device()
B, T, V = 2, 4, 16
logits = torch.randn(B, T + 1, V, device=device, dtype=torch.float32).contiguous()
completion_ids = torch.randint(0, V, (B, T), device=device)
completion_mask = torch.randint(0, 2, (B, T), device=device, dtype=torch.long)
completion_mask[:, 0] = 1 # ensure each sequence has at least one valid token
advantages = torch.randn(B, device=device, dtype=torch.float32)
old_logp = torch.randn(B, T, device=device, dtype=torch.float32)
ref_logp = torch.randn(B, T, device=device, dtype=torch.float32) if beta != 0.0 else None
per_token_loss, per_token_kl, is_clipped = triton_grpo_loss(
logits=logits,
old_logp=old_logp,
ref_logp=ref_logp,
completion_ids=completion_ids,
advantages=advantages,
completion_mask=completion_mask,
temperature=1.0,
beta=beta,
eps_low=0.2,
eps_high=0.2,
inplace=False,
loss_type=loss_type,
max_completion_length=T,
reduce=False,
)
logits_main = logits[:, :-1, :]
log_probs = torch.log_softmax(logits_main, dim=-1)
per_token_logps = log_probs.gather(dim=-1, index=completion_ids.unsqueeze(-1)).squeeze(-1)
ref_tokens = ref_logp if ref_logp is not None else per_token_logps.detach()
reference_loss, reference_kl, reference_is_clipped = TorchLMHeadGRPO.compute_per_token_components(
per_token_logps,
completion_mask.float(),
advantages,
old_logp,
ref_tokens,
0.2,
0.2,
beta,
"token",
)
mask = completion_mask.float()
mask_bool = mask.bool()
assert_verbose_allclose(per_token_loss, reference_loss * mask)
assert torch.equal(is_clipped.bool()[mask_bool], reference_is_clipped[mask_bool])
if beta != 0.0:
assert_verbose_allclose(per_token_kl, reference_kl * mask)
else:
assert per_token_kl is None
reduced_loss, metrics = triton_grpo_loss(
logits=logits,
old_logp=old_logp,
ref_logp=ref_logp,
completion_ids=completion_ids,
advantages=advantages,
completion_mask=completion_mask,
temperature=1.0,
beta=beta,
eps_low=0.2,
eps_high=0.2,
inplace=False,
loss_type=loss_type,
max_completion_length=T,
reduce=True,
)
expected_loss = _reduce_grpo_loss(reference_loss, completion_mask, loss_type, T)
assert_verbose_allclose(reduced_loss, expected_loss)
if beta != 0.0:
assert_verbose_allclose(metrics[0], _masked_mean(reference_kl, completion_mask))
clip_metric = metrics[1]
else:
clip_metric = metrics[0]
assert_verbose_allclose(clip_metric, _masked_mean(reference_is_clipped.float(), completion_mask))
def _reference_per_token_loss(
logits,
completion_ids,
completion_mask,
advantages,
old_logp,
ref_logp,
beta,
eps_low,
eps_high,
temperature=1.0,
):
logits = logits[:, :-1, :] / temperature
log_probs = torch.log_softmax(logits, dim=-1)
per_token_logps = log_probs.gather(-1, completion_ids.unsqueeze(-1)).squeeze(-1)
old = old_logp if old_logp is not None else per_token_logps.detach()
coef_1 = torch.exp(per_token_logps - old)
coef_2 = torch.clamp(coef_1, 1 - eps_low, 1 + eps_high)
per_token_loss1 = coef_1 * advantages.unsqueeze(1)
per_token_loss2 = coef_2 * advantages.unsqueeze(1)
per_token_loss = -torch.minimum(per_token_loss1, per_token_loss2)
is_clipped = per_token_loss1 < per_token_loss2
mask = completion_mask.to(torch.bool)
per_token_loss = per_token_loss.masked_fill(~mask, 0.0)
is_clipped = is_clipped & mask
if beta != 0.0:
kl = torch.exp(ref_logp - per_token_logps) - (ref_logp - per_token_logps) - 1.0
kl = kl.masked_fill(~mask, 0.0)
per_token_loss = per_token_loss + beta * kl
else:
kl = None
return {
"per_token_loss": per_token_loss,
"kl": kl,
"is_clipped": is_clipped,
}
def _masked_mean(values, mask):
mask = mask.to(values.dtype)
return (values * mask).sum() / mask.sum().clamp(min=1.0)
+174
View File
@@ -30,6 +30,8 @@ from liger_kernel.transformers import apply_liger_kernel_to_glm4
from liger_kernel.transformers import apply_liger_kernel_to_glm4v
from liger_kernel.transformers import apply_liger_kernel_to_glm4v_moe
from liger_kernel.transformers import apply_liger_kernel_to_granite
from liger_kernel.transformers import apply_liger_kernel_to_hunyuan_v1_dense
from liger_kernel.transformers import apply_liger_kernel_to_hunyuan_v1_moe
from liger_kernel.transformers import apply_liger_kernel_to_internvl
from liger_kernel.transformers import apply_liger_kernel_to_llama
from liger_kernel.transformers import apply_liger_kernel_to_llama4
@@ -38,6 +40,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_mistral
from liger_kernel.transformers import apply_liger_kernel_to_mixtral
from liger_kernel.transformers import apply_liger_kernel_to_mllama
from liger_kernel.transformers import apply_liger_kernel_to_olmo2
from liger_kernel.transformers import apply_liger_kernel_to_olmo3
from liger_kernel.transformers import apply_liger_kernel_to_phi3
from liger_kernel.transformers import apply_liger_kernel_to_qwen2
from liger_kernel.transformers import apply_liger_kernel_to_qwen2_5_vl
@@ -62,6 +65,8 @@ from test.utils import revert_liger_kernel_to_glm4
from test.utils import revert_liger_kernel_to_glm4v
from test.utils import revert_liger_kernel_to_glm4v_moe
from test.utils import revert_liger_kernel_to_granite
from test.utils import revert_liger_kernel_to_hunyuan_v1
from test.utils import revert_liger_kernel_to_hunyuan_v1_moe
from test.utils import revert_liger_kernel_to_internvl
from test.utils import revert_liger_kernel_to_llama
from test.utils import revert_liger_kernel_to_llama4
@@ -70,6 +75,7 @@ from test.utils import revert_liger_kernel_to_mistral
from test.utils import revert_liger_kernel_to_mixtral
from test.utils import revert_liger_kernel_to_mllama
from test.utils import revert_liger_kernel_to_olmo2
from test.utils import revert_liger_kernel_to_olmo3
from test.utils import revert_liger_kernel_to_phi3
from test.utils import revert_liger_kernel_to_qwen2
from test.utils import revert_liger_kernel_to_qwen2_5_vl
@@ -190,6 +196,15 @@ try:
except ImportError:
OLMO2_AVAILABLE = False
try:
# OLMO3 is only available in transformers>=4.57.0
from transformers.models.olmo3.configuration_olmo3 import Olmo3Config
from transformers.models.olmo3.modeling_olmo3 import Olmo3ForCausalLM
OLMO3_AVAILABLE = True
except ImportError:
OLMO3_AVAILABLE = False
try:
# Glm4 is only available in transformers>=4.51.3
from transformers.models.glm4.configuration_glm4 import Glm4Config
@@ -261,6 +276,16 @@ try:
except ImportError:
QWEN3NEXT_AVAILABLE = False
try:
from transformers.models.hunyuan_v1_dense.configuration_hunyuan_v1_dense import HunYuanDenseV1Config
from transformers.models.hunyuan_v1_dense.modeling_hunyuan_v1_dense import HunYuanDenseV1ForCausalLM
from transformers.models.hunyuan_v1_moe.configuration_hunyuan_v1_moe import HunYuanMoEV1Config
from transformers.models.hunyuan_v1_moe.modeling_hunyuan_v1_moe import HunYuanMoEV1ForCausalLM
HUNYUAN_V1_AVAILABLE = True
except ImportError:
HUNYUAN_V1_AVAILABLE = False
from liger_kernel.utils import infer_device
device = infer_device()
@@ -995,6 +1020,35 @@ if OLMO2_AVAILABLE:
),
)
if OLMO3_AVAILABLE:
MINI_MODEL_SETUPS["mini_olmo3"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_olmo3,
liger_kernel_patch_revert_func=revert_liger_kernel_to_olmo3,
model_class=Olmo3ForCausalLM,
mini_model_config=Olmo3Config(
bos_token_id=1, # 128000
eos_token_id=2, # 128001
pad_token_id=2,
cross_attention_layers=None,
dropout=0,
hidden_act="silu",
hidden_size=1024, # 4096
initializer_range=0.02,
intermediate_size=2048, # 14336
max_position_embeddings=4096,
num_attention_heads=8, # 32
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
attn_implementation="sdpa", # default value, pytorch native attention
),
)
if GLM4_AVAILABLE:
MINI_MODEL_SETUPS["mini_glm4"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_glm4,
@@ -1304,6 +1358,69 @@ if QWEN3NEXT_AVAILABLE:
)
if HUNYUAN_V1_AVAILABLE:
MINI_MODEL_SETUPS["mini_hunyuan_v1"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_hunyuan_v1_dense,
liger_kernel_patch_revert_func=revert_liger_kernel_to_hunyuan_v1,
model_class=HunYuanDenseV1ForCausalLM,
mini_model_config=HunYuanDenseV1Config(
attention_dropout=0.0,
bos_token_id=1,
eos_token_id=2,
hidden_act="silu",
num_hidden_layers=4,
hidden_size=896,
intermediate_size=4864,
num_attention_heads=8,
head_dim=112,
rms_norm_eps=1e-6,
tie_word_embeddings=True,
max_position_embeddings=32768,
initializer_range=0.02,
norm_eps=1e-6,
num_key_value_heads=2,
rope_theta=10000.0,
partial_rotary_factor=1.0,
vocab_size=32000,
use_cache=True,
attn_implementation="sdpa",
),
)
MINI_MODEL_SETUPS["mini_hunyuan_v1_moe"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_hunyuan_v1_moe,
liger_kernel_patch_revert_func=revert_liger_kernel_to_hunyuan_v1_moe,
model_class=HunYuanMoEV1ForCausalLM,
mini_model_config=HunYuanMoEV1Config(
vocab_size=32000,
hidden_size=128,
intermediate_size=512,
head_dim=16,
num_hidden_layers=2,
num_attention_heads=8,
num_key_value_heads=2,
hidden_act="silu",
max_position_embeddings=32768,
initializer_range=0.02,
rms_norm_eps=1e-5,
use_cache=True,
pad_token_id=0,
bos_token_id=1,
eos_token_id=2,
eod_token_id=3,
sep_token_id=4,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
attention_dropout=0.0,
num_experts=2,
moe_topk=1,
attn_implementation="sdpa",
),
)
def create_model(model_name="mini_llama4"):
"""
Create a mini version model
@@ -1674,6 +1791,25 @@ def run_mini_model(
),
],
),
pytest.param(
"mini_olmo3",
32,
1e-5,
torch.bfloat16,
1e-2,
1e-2,
1e-1,
1e-2,
1e-2,
1e-2,
marks=[
pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
pytest.mark.skipif(
not OLMO3_AVAILABLE,
reason="OLMO3 not available in this version of transformers",
),
],
),
pytest.param(
"mini_glm4",
32,
@@ -1885,6 +2021,44 @@ def run_mini_model(
),
],
),
pytest.param(
"mini_hunyuan_v1",
32,
1e-5,
torch.bfloat16,
1e-2,
5e-2,
1e-1,
1e-2,
1e-2,
1e-2,
marks=[
pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
pytest.mark.skipif(
not HUNYUAN_V1_AVAILABLE,
reason="Hunyuan_v1 not available in this version of transformers",
),
],
),
pytest.param(
"mini_hunyuan_v1_moe",
32,
1e-5,
torch.bfloat16,
5e-2,
5e-2,
1e-1, # 1e-1
1e-1, # 1e-2
1e-2,
1e-2,
marks=[
pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
pytest.mark.skipif(
not HUNYUAN_V1_AVAILABLE,
reason="Hunyuan_v1_moe not available in this version of transformers",
),
],
),
],
)
def test_mini_model(
@@ -30,6 +30,8 @@ from liger_kernel.transformers import apply_liger_kernel_to_glm4
from liger_kernel.transformers import apply_liger_kernel_to_glm4v
from liger_kernel.transformers import apply_liger_kernel_to_glm4v_moe
from liger_kernel.transformers import apply_liger_kernel_to_granite
from liger_kernel.transformers import apply_liger_kernel_to_hunyuan_v1_dense
from liger_kernel.transformers import apply_liger_kernel_to_hunyuan_v1_moe
from liger_kernel.transformers import apply_liger_kernel_to_internvl
from liger_kernel.transformers import apply_liger_kernel_to_llama
from liger_kernel.transformers import apply_liger_kernel_to_llama4
@@ -38,6 +40,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_mistral
from liger_kernel.transformers import apply_liger_kernel_to_mixtral
from liger_kernel.transformers import apply_liger_kernel_to_mllama
from liger_kernel.transformers import apply_liger_kernel_to_olmo2
from liger_kernel.transformers import apply_liger_kernel_to_olmo3
from liger_kernel.transformers import apply_liger_kernel_to_phi3
from liger_kernel.transformers import apply_liger_kernel_to_qwen2
from liger_kernel.transformers import apply_liger_kernel_to_qwen2_5_vl
@@ -62,6 +65,8 @@ from test.utils import revert_liger_kernel_to_glm4
from test.utils import revert_liger_kernel_to_glm4v
from test.utils import revert_liger_kernel_to_glm4v_moe
from test.utils import revert_liger_kernel_to_granite
from test.utils import revert_liger_kernel_to_hunyuan_v1
from test.utils import revert_liger_kernel_to_hunyuan_v1_moe
from test.utils import revert_liger_kernel_to_internvl
from test.utils import revert_liger_kernel_to_llama
from test.utils import revert_liger_kernel_to_llama4
@@ -70,6 +75,7 @@ from test.utils import revert_liger_kernel_to_mistral
from test.utils import revert_liger_kernel_to_mixtral
from test.utils import revert_liger_kernel_to_mllama
from test.utils import revert_liger_kernel_to_olmo2
from test.utils import revert_liger_kernel_to_olmo3
from test.utils import revert_liger_kernel_to_phi3
from test.utils import revert_liger_kernel_to_qwen2
from test.utils import revert_liger_kernel_to_qwen2_5_vl
@@ -182,6 +188,15 @@ try:
except ImportError:
OLMO2_AVAILABLE = False
try:
# OLMO3 is only available in transformers>=4.57.0
from transformers.models.olmo3.configuration_olmo3 import Olmo3Config
from transformers.models.olmo3.modeling_olmo3 import Olmo3ForCausalLM
OLMO3_AVAILABLE = True
except ImportError:
OLMO3_AVAILABLE = False
try:
# Glm4 is only available in transformers>=4.51.3
from transformers.models.glm4.configuration_glm4 import Glm4Config
@@ -254,6 +269,17 @@ except ImportError:
QWEN3NEXT_AVAILABLE = False
try:
from transformers.models.hunyuan_v1_dense.configuration_hunyuan_v1_dense import HunYuanDenseV1Config
from transformers.models.hunyuan_v1_dense.modeling_hunyuan_v1_dense import HunYuanDenseV1ForCausalLM
from transformers.models.hunyuan_v1_moe.configuration_hunyuan_v1_moe import HunYuanMoEV1Config
from transformers.models.hunyuan_v1_moe.modeling_hunyuan_v1_moe import HunYuanMoEV1ForCausalLM
HUNYUAN_V1_AVAILABLE = True
except ImportError:
HUNYUAN_V1_AVAILABLE = False
from liger_kernel.utils import infer_device
device = infer_device()
@@ -990,6 +1016,35 @@ if OLMO2_AVAILABLE:
),
)
if OLMO3_AVAILABLE:
MINI_MODEL_SETUPS["mini_olmo3"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_olmo3,
liger_kernel_patch_revert_func=revert_liger_kernel_to_olmo3,
model_class=Olmo3ForCausalLM,
mini_model_config=Olmo3Config(
bos_token_id=1, # 128000
eos_token_id=2, # 128001
pad_token_id=2,
cross_attention_layers=None,
dropout=0,
hidden_act="silu",
hidden_size=1024, # 4096
initializer_range=0.02,
intermediate_size=2048, # 14336
max_position_embeddings=4096,
num_attention_heads=8, # 32
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
attn_implementation="sdpa", # default value, pytorch native attention
),
)
if GLM4_AVAILABLE:
MINI_MODEL_SETUPS["mini_glm4"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_glm4,
@@ -1298,6 +1353,69 @@ if QWEN3NEXT_AVAILABLE:
)
if HUNYUAN_V1_AVAILABLE:
MINI_MODEL_SETUPS["mini_hunyuan_v1"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_hunyuan_v1_dense,
liger_kernel_patch_revert_func=revert_liger_kernel_to_hunyuan_v1,
model_class=HunYuanDenseV1ForCausalLM,
mini_model_config=HunYuanDenseV1Config(
attention_dropout=0.0,
bos_token_id=1,
eos_token_id=2,
hidden_act="silu",
num_hidden_layers=4,
hidden_size=896,
intermediate_size=4864,
num_attention_heads=8,
head_dim=112,
rms_norm_eps=1e-6,
tie_word_embeddings=True,
max_position_embeddings=32768,
initializer_range=0.02,
norm_eps=1e-6,
num_key_value_heads=2,
rope_theta=10000.0,
partial_rotary_factor=1.0,
vocab_size=32000,
use_cache=True,
attn_implementation="sdpa",
),
)
MINI_MODEL_SETUPS["mini_hunyuan_v1_moe"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_hunyuan_v1_moe,
liger_kernel_patch_revert_func=revert_liger_kernel_to_hunyuan_v1_moe,
model_class=HunYuanMoEV1ForCausalLM,
mini_model_config=HunYuanMoEV1Config(
vocab_size=32000,
hidden_size=128,
intermediate_size=512,
head_dim=16,
num_hidden_layers=2,
num_attention_heads=8,
num_key_value_heads=2,
hidden_act="silu",
max_position_embeddings=32768,
initializer_range=0.02,
rms_norm_eps=1e-5,
use_cache=True,
pad_token_id=0,
bos_token_id=1,
eos_token_id=2,
eod_token_id=3,
sep_token_id=4,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
attention_dropout=0.0,
num_experts=2,
moe_topk=1,
attn_implementation="sdpa",
),
)
def create_model(model_name="mini_llama3"):
"""
Create a mini version model
@@ -1690,6 +1808,25 @@ def run_mini_model(
),
],
),
pytest.param(
"mini_olmo3",
32,
1e-4,
torch.bfloat16,
1e-2,
1e-2,
1e-1,
1e-2,
1e-2,
1e-2,
marks=[
pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
pytest.mark.skipif(
not OLMO3_AVAILABLE,
reason="OLMO3 not available in this version of transformers",
),
],
),
pytest.param(
"mini_glm4",
32,
@@ -1858,6 +1995,44 @@ def run_mini_model(
),
],
),
pytest.param(
"mini_hunyuan_v1",
32,
1e-5,
torch.bfloat16,
1e-2,
5e-2,
1e-1,
1e-2,
1e-2,
1e-2,
marks=[
pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
pytest.mark.skipif(
not HUNYUAN_V1_AVAILABLE,
reason="Hunyuan_v1 not available in this version of transformers",
),
],
),
pytest.param(
"mini_hunyuan_v1_moe",
32,
1e-5,
torch.bfloat16,
1e-2,
5e-2,
1e-1,
1e-2,
1e-2,
1e-2,
marks=[
pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
pytest.mark.skipif(
not HUNYUAN_V1_AVAILABLE,
reason="Hunyuan_v1_moe not available in this version of transformers",
),
],
),
],
)
def test_mini_model(
+159
View File
@@ -30,6 +30,8 @@ from liger_kernel.transformers import apply_liger_kernel_to_glm4
from liger_kernel.transformers import apply_liger_kernel_to_glm4v
from liger_kernel.transformers import apply_liger_kernel_to_glm4v_moe
from liger_kernel.transformers import apply_liger_kernel_to_granite
from liger_kernel.transformers import apply_liger_kernel_to_hunyuan_v1_dense
from liger_kernel.transformers import apply_liger_kernel_to_hunyuan_v1_moe
from liger_kernel.transformers import apply_liger_kernel_to_internvl
from liger_kernel.transformers import apply_liger_kernel_to_llama
from liger_kernel.transformers import apply_liger_kernel_to_llama4
@@ -38,6 +40,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_mistral
from liger_kernel.transformers import apply_liger_kernel_to_mixtral
from liger_kernel.transformers import apply_liger_kernel_to_mllama
from liger_kernel.transformers import apply_liger_kernel_to_olmo2
from liger_kernel.transformers import apply_liger_kernel_to_olmo3
from liger_kernel.transformers import apply_liger_kernel_to_phi3
from liger_kernel.transformers import apply_liger_kernel_to_qwen2
from liger_kernel.transformers import apply_liger_kernel_to_qwen2_5_vl
@@ -62,6 +65,8 @@ from test.utils import revert_liger_kernel_to_glm4
from test.utils import revert_liger_kernel_to_glm4v
from test.utils import revert_liger_kernel_to_glm4v_moe
from test.utils import revert_liger_kernel_to_granite
from test.utils import revert_liger_kernel_to_hunyuan_v1
from test.utils import revert_liger_kernel_to_hunyuan_v1_moe
from test.utils import revert_liger_kernel_to_internvl
from test.utils import revert_liger_kernel_to_llama
from test.utils import revert_liger_kernel_to_llama4
@@ -70,6 +75,7 @@ from test.utils import revert_liger_kernel_to_mistral
from test.utils import revert_liger_kernel_to_mixtral
from test.utils import revert_liger_kernel_to_mllama
from test.utils import revert_liger_kernel_to_olmo2
from test.utils import revert_liger_kernel_to_olmo3
from test.utils import revert_liger_kernel_to_phi3
from test.utils import revert_liger_kernel_to_qwen2
from test.utils import revert_liger_kernel_to_qwen2_5_vl
@@ -178,6 +184,15 @@ try:
except ImportError:
OLMO2_AVAILABLE = False
try:
# OLMO3 is only available in transformers>=4.57.0
from transformers.models.olmo3.configuration_olmo3 import Olmo3Config
from transformers.models.olmo3.modeling_olmo3 import Olmo3ForCausalLM
OLMO3_AVAILABLE = True
except ImportError:
OLMO3_AVAILABLE = False
try:
# Glm4 is only available in transformers>=4.51.3
from transformers.models.glm4.configuration_glm4 import Glm4Config
@@ -259,6 +274,16 @@ try:
except ImportError:
QWEN3NEXT_AVAILABLE = False
try:
from transformers.models.hunyuan_v1_dense.configuration_hunyuan_v1_dense import HunYuanDenseV1Config
from transformers.models.hunyuan_v1_dense.modeling_hunyuan_v1_dense import HunYuanDenseV1ForCausalLM
from transformers.models.hunyuan_v1_moe.configuration_hunyuan_v1_moe import HunYuanMoEV1Config
from transformers.models.hunyuan_v1_moe.modeling_hunyuan_v1_moe import HunYuanMoEV1ForCausalLM
HUNYUAN_V1_AVAILABLE = True
except ImportError:
HUNYUAN_V1_AVAILABLE = False
from liger_kernel.utils import infer_device
device = infer_device()
@@ -931,6 +956,35 @@ if OLMO2_AVAILABLE:
),
)
if OLMO3_AVAILABLE:
MINI_MODEL_SETUPS["mini_olmo3"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_olmo3,
liger_kernel_patch_revert_func=revert_liger_kernel_to_olmo3,
model_class=Olmo3ForCausalLM,
mini_model_config=Olmo3Config(
bos_token_id=1, # 128000
eos_token_id=2, # 128001
pad_token_id=2,
cross_attention_layers=None,
dropout=0,
hidden_act="silu",
hidden_size=1024, # 4096
initializer_range=0.02,
intermediate_size=2048, # 14336
max_position_embeddings=4096,
num_attention_heads=8, # 32
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
attn_implementation="sdpa", # default value, pytorch native attention
),
)
if GLM4_AVAILABLE:
MINI_MODEL_SETUPS["mini_glm4"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_glm4,
@@ -1296,6 +1350,63 @@ if QWEN3NEXT_AVAILABLE:
),
)
if HUNYUAN_V1_AVAILABLE:
MINI_MODEL_SETUPS["mini_hunyuan_v1"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_hunyuan_v1_dense,
liger_kernel_patch_revert_func=revert_liger_kernel_to_hunyuan_v1,
model_class=HunYuanDenseV1ForCausalLM,
mini_model_config=HunYuanDenseV1Config(
attention_dropout=0.0,
bos_token_id=1,
eos_token_id=2,
hidden_act="silu",
num_hidden_layers=4,
hidden_size=896,
intermediate_size=4864,
num_attention_heads=8,
head_dim=112,
rms_norm_eps=1e-6,
tie_word_embeddings=True,
max_position_embeddings=32768,
initializer_range=0.02,
norm_eps=1e-6,
num_key_value_heads=2,
rope_theta=10000.0,
partial_rotary_factor=1.0,
vocab_size=32000,
use_cache=True,
attn_implementation="sdpa",
),
)
MINI_MODEL_SETUPS["mini_hunyuan_v1_moe"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_hunyuan_v1_moe,
liger_kernel_patch_revert_func=revert_liger_kernel_to_hunyuan_v1_moe,
model_class=HunYuanMoEV1ForCausalLM,
mini_model_config=HunYuanMoEV1Config(
hidden_act="silu",
attention_dropout=0.0,
num_hidden_layers=4,
hidden_size=896,
intermediate_size=4864,
num_attention_heads=8,
head_dim=112,
rms_norm_eps=1e-6,
tie_word_embeddings=True,
max_position_embeddings=32768,
initializer_range=0.02,
norm_eps=1e-6,
num_key_value_heads=2,
rope_theta=10000.0,
partial_rotary_factor=1.0,
vocab_size=32000,
num_experts=8,
moe_topk=2,
use_cache=True,
attn_implementation="sdpa",
),
)
def create_model(model_name="mini_llama3"):
"""
@@ -1578,6 +1689,22 @@ def run_mini_model(
reason="OLMO2 not available in this version of transformers",
),
),
pytest.param(
"mini_olmo3",
32,
1e-4,
torch.float32,
1e-8,
1e-5,
5e-3,
1e-5,
5e-3,
1e-5,
marks=pytest.mark.skipif(
not OLMO3_AVAILABLE,
reason="OLMO3 not available in this version of transformers",
),
),
pytest.param(
"mini_glm4",
32,
@@ -1737,6 +1864,38 @@ def run_mini_model(
),
],
),
pytest.param(
"mini_hunyuan_v1",
32,
1e-4,
torch.float32,
1e-8,
1e-5,
5e-3,
1e-5,
5e-3,
1e-5,
marks=pytest.mark.skipif(
not HUNYUAN_V1_AVAILABLE,
reason="Hunyuan_v1 not available in this version of transformers",
),
),
pytest.param(
"mini_hunyuan_v1_moe",
32,
1e-5,
torch.float32,
1e-8,
1e-5,
5e-3,
1e-5,
5e-3,
1e-5,
marks=pytest.mark.skipif(
not HUNYUAN_V1_AVAILABLE,
reason="Hunyuan_v1_moe not available in this version of transformers",
),
),
],
)
def test_mini_model(
@@ -30,6 +30,8 @@ from liger_kernel.transformers import apply_liger_kernel_to_glm4
from liger_kernel.transformers import apply_liger_kernel_to_glm4v
from liger_kernel.transformers import apply_liger_kernel_to_glm4v_moe
from liger_kernel.transformers import apply_liger_kernel_to_granite
from liger_kernel.transformers import apply_liger_kernel_to_hunyuan_v1_dense
from liger_kernel.transformers import apply_liger_kernel_to_hunyuan_v1_moe
from liger_kernel.transformers import apply_liger_kernel_to_internvl
from liger_kernel.transformers import apply_liger_kernel_to_llama
from liger_kernel.transformers import apply_liger_kernel_to_llama4
@@ -38,6 +40,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_mistral
from liger_kernel.transformers import apply_liger_kernel_to_mixtral
from liger_kernel.transformers import apply_liger_kernel_to_mllama
from liger_kernel.transformers import apply_liger_kernel_to_olmo2
from liger_kernel.transformers import apply_liger_kernel_to_olmo3
from liger_kernel.transformers import apply_liger_kernel_to_phi3
from liger_kernel.transformers import apply_liger_kernel_to_qwen2
from liger_kernel.transformers import apply_liger_kernel_to_qwen2_5_vl
@@ -62,6 +65,8 @@ from test.utils import revert_liger_kernel_to_glm4
from test.utils import revert_liger_kernel_to_glm4v
from test.utils import revert_liger_kernel_to_glm4v_moe
from test.utils import revert_liger_kernel_to_granite
from test.utils import revert_liger_kernel_to_hunyuan_v1
from test.utils import revert_liger_kernel_to_hunyuan_v1_moe
from test.utils import revert_liger_kernel_to_internvl
from test.utils import revert_liger_kernel_to_llama
from test.utils import revert_liger_kernel_to_llama4
@@ -70,6 +75,7 @@ from test.utils import revert_liger_kernel_to_mistral
from test.utils import revert_liger_kernel_to_mixtral
from test.utils import revert_liger_kernel_to_mllama
from test.utils import revert_liger_kernel_to_olmo2
from test.utils import revert_liger_kernel_to_olmo3
from test.utils import revert_liger_kernel_to_phi3
from test.utils import revert_liger_kernel_to_qwen2
from test.utils import revert_liger_kernel_to_qwen2_5_vl
@@ -202,6 +208,15 @@ try:
except ImportError:
OLMO2_AVAILABLE = False
try:
# OLMO3 is only available in transformers>=4.57.0
from transformers.models.olmo3.configuration_olmo3 import Olmo3Config
from transformers.models.olmo3.modeling_olmo3 import Olmo3ForCausalLM
OLMO3_AVAILABLE = True
except ImportError:
OLMO3_AVAILABLE = False
try:
# Glm4 is only available in transformers>=4.51.3
from transformers.models.glm4.configuration_glm4 import Glm4Config
@@ -273,6 +288,16 @@ try:
except ImportError:
QWEN3NEXT_AVAILABLE = False
try:
from transformers.models.hunyuan_v1_dense.configuration_hunyuan_v1_dense import HunYuanDenseV1Config
from transformers.models.hunyuan_v1_dense.modeling_hunyuan_v1_dense import HunYuanDenseV1ForCausalLM
from transformers.models.hunyuan_v1_moe.configuration_hunyuan_v1_moe import HunYuanMoEV1Config
from transformers.models.hunyuan_v1_moe.modeling_hunyuan_v1_moe import HunYuanMoEV1ForCausalLM
HUNYUAN_V1_AVAILABLE = True
except ImportError:
HUNYUAN_V1_AVAILABLE = False
from liger_kernel.utils import infer_device
device = infer_device()
@@ -1008,6 +1033,35 @@ if OLMO2_AVAILABLE:
),
)
if OLMO3_AVAILABLE:
MINI_MODEL_SETUPS["mini_olmo3"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_olmo3,
liger_kernel_patch_revert_func=revert_liger_kernel_to_olmo3,
model_class=Olmo3ForCausalLM,
mini_model_config=Olmo3Config(
bos_token_id=1, # 128000
eos_token_id=2, # 128001
pad_token_id=2,
cross_attention_layers=None,
dropout=0,
hidden_act="silu",
hidden_size=1024, # 4096
initializer_range=0.02,
intermediate_size=2048, # 14336
max_position_embeddings=4096,
num_attention_heads=8, # 32
num_hidden_layers=4, # 40
num_key_value_heads=2, # 8
rms_norm_eps=1e-5,
rope_scaling=None,
rope_theta=500_000,
tie_word_embeddings=False,
use_cache=True,
vocab_size=32000, # 128256,
attn_implementation="sdpa", # default value, pytorch native attention
),
)
if GLM4_AVAILABLE:
MINI_MODEL_SETUPS["mini_glm4"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_glm4,
@@ -1316,6 +1370,64 @@ if QWEN3NEXT_AVAILABLE:
)
if HUNYUAN_V1_AVAILABLE:
MINI_MODEL_SETUPS["mini_hunyuan_v1"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_hunyuan_v1_dense,
liger_kernel_patch_revert_func=revert_liger_kernel_to_hunyuan_v1,
model_class=HunYuanDenseV1ForCausalLM,
mini_model_config=HunYuanDenseV1Config(
attention_dropout=0.0,
bos_token_id=1,
eos_token_id=2,
hidden_act="silu",
num_hidden_layers=4,
hidden_size=896,
intermediate_size=4864,
num_attention_heads=8,
head_dim=112,
rms_norm_eps=1e-6,
tie_word_embeddings=True,
max_position_embeddings=32768,
initializer_range=0.02,
norm_eps=1e-6,
num_key_value_heads=2,
rope_theta=10000.0,
partial_rotary_factor=1.0,
vocab_size=32000,
use_cache=True,
attn_implementation="sdpa",
),
)
MINI_MODEL_SETUPS["mini_hunyuan_v1_moe"] = MiniModelConfig(
liger_kernel_patch_func=apply_liger_kernel_to_hunyuan_v1_moe,
liger_kernel_patch_revert_func=revert_liger_kernel_to_hunyuan_v1_moe,
model_class=HunYuanMoEV1ForCausalLM,
mini_model_config=HunYuanMoEV1Config(
hidden_act="silu",
attention_dropout=0.0,
num_hidden_layers=4,
hidden_size=896,
intermediate_size=4864,
num_attention_heads=8,
head_dim=112,
rms_norm_eps=1e-6,
tie_word_embeddings=True,
max_position_embeddings=32768,
initializer_range=0.02,
norm_eps=1e-6,
num_key_value_heads=2,
rope_theta=10000.0,
partial_rotary_factor=1.0,
vocab_size=32000,
num_experts=8,
moe_topk=2,
use_cache=True,
attn_implementation="sdpa",
),
)
def create_model(model_name="mini_llama3"):
"""
Create a mini version model
@@ -1581,6 +1693,22 @@ def run_mini_model(
reason="OLMO2 not available in this version of transformers",
),
),
pytest.param(
"mini_olmo3",
32,
1e-4,
torch.float32,
1e-8,
1e-5,
5e-3,
1e-5,
5e-3,
1e-5,
marks=pytest.mark.skipif(
not OLMO3_AVAILABLE,
reason="OLMO3 not available in this version of transformers",
),
),
pytest.param(
"mini_glm4",
32,
@@ -1723,6 +1851,42 @@ def run_mini_model(
),
],
),
pytest.param(
"mini_hunyuan_v1",
32,
1e-5,
torch.float32,
1e-2,
5e-2,
1e-1,
1e-2,
1e-2,
1e-2,
marks=[
pytest.mark.skipif(
not HUNYUAN_V1_AVAILABLE,
reason="Hunyuan_v1 not available in this version of transformers",
),
],
),
pytest.param(
"mini_hunyuan_v1_moe",
32,
1e-5,
torch.float32,
1e-2,
5e-2,
1e-1,
1e-2,
1e-2,
1e-2,
marks=[
pytest.mark.skipif(
not HUNYUAN_V1_AVAILABLE,
reason="Hunyuan_v1_moe not available in this version of transformers",
),
],
),
],
)
def test_mini_model(
+1
View File
@@ -68,6 +68,7 @@ def test_liger_layer_norm(
[
(2, 8, 64),
(4, 16, 128),
(3, 512, 128),
],
)
@pytest.mark.parametrize(
+151
View File
@@ -143,6 +143,15 @@ def is_olmo2_available():
return False
def is_olmo3_available():
try:
import transformers.models.olmo3 # noqa: F401
return True
except ImportError:
return False
def is_glm4_available():
try:
import transformers.models.glm4 # noqa: F401
@@ -206,6 +215,15 @@ def is_qwen3_next_available():
return False
def is_hunyuan_v1_available():
try:
import transformers.models.hunyuan_v1_dense # noqa: F401
return True
except ImportError:
return False
def test_import_from_root():
try:
from liger_kernel.transformers import AutoLigerKernelForCausalLM # noqa: F401
@@ -2469,6 +2487,52 @@ def test_apply_liger_kernel_to_instance_for_olmo2():
pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
@pytest.mark.skipif(not is_olmo3_available(), reason="olmo3 module not available")
def test_apply_liger_kernel_to_instance_for_olmo3():
# Ensure any monkey patching is cleaned up for subsequent tests
with patch("transformers.models.olmo3.modeling_olmo3"):
from liger_kernel.transformers.model.olmo3 import lce_forward as olmo3_lce_forward
# Instantiate a dummy model
config = transformers.models.olmo3.configuration_olmo3.Olmo3Config(
dtype=torch.bfloat16,
rms_norm_eps=1e-5,
hidden_size=32,
intermediate_size=64,
hidden_act="silu",
num_hidden_layers=2,
)
dummy_model_instance = AutoModelForCausalLM.from_config(config)
# Check that model instance variables are not yet patched with Liger modules
assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(olmo3_lce_forward)
assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerSwiGLUMLP.forward)
assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(layer.post_feedforward_layernorm.forward) != inspect.getsource(
LigerRMSNorm.forward
)
# Test applying kernels to the model instance
_apply_liger_kernel_to_instance(model=dummy_model_instance)
# Check that the model's instance variables were correctly patched with Liger modules
assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(olmo3_lce_forward)
assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerSwiGLUMLP.forward)
assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(layer.post_feedforward_layernorm.forward) == inspect.getsource(
LigerRMSNorm.forward
)
try:
print(dummy_model_instance)
except Exception as e:
pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
@pytest.mark.skipif(not is_glm4_available(), reason="glm4 module not available")
def test_apply_liger_kernel_to_instance_for_glm4():
# Ensure any monkey patching is cleaned up for subsequent tests
@@ -2790,3 +2854,90 @@ def test_apply_liger_kernel_to_instance_for_qwen3_next():
print(dummy_model_instance)
except Exception as e:
pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
@pytest.mark.skipif(not is_hunyuan_v1_available(), reason="hunyuan_v1 module not available")
def test_apply_liger_kernel_to_instance_for_hunyuan_v1_moe():
# Ensure any monkey patching is cleaned up for subsequent tests
with patch("transformers.models.hunyuan_v1_moe.modeling_hunyuan_v1_moe"):
from liger_kernel.transformers.model.hunyuan_v1 import lce_forward as hunyuan_v1_moe_lce_forward
# Instantiate a dummy model
config = transformers.models.hunyuan_v1_moe.configuration_hunyuan_v1_moe.HunYuanMoEV1Config(
torch_dtype=torch.bfloat16,
rms_norm_eps=1e-5,
hidden_size=32,
intermediate_size=64,
hidden_act="silu",
num_hidden_layers=2,
head_dim=1,
)
dummy_model_instance = AutoModelForCausalLM.from_config(config)
# Check that model instance variables are not yet patched with Liger modules
assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(hunyuan_v1_moe_lce_forward)
assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward)
assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
# Test applying kernels to the model instance
_apply_liger_kernel_to_instance(model=dummy_model_instance)
# Check that the model's instance variables were correctly patched with Liger modules
assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(hunyuan_v1_moe_lce_forward)
assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
for mlp_expert in layer.mlp.experts:
assert inspect.getsource(mlp_expert.forward) == inspect.getsource(LigerQwen3MoeSwiGLUMLP.forward)
assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
try:
print(dummy_model_instance)
except Exception as e:
pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
@pytest.mark.skipif(not is_hunyuan_v1_available(), reason="hunyuan_v1_dense module not available")
def test_apply_liger_kernel_to_instance_for_hunyuan_v1_dense():
# Ensure any monkey patching is cleaned up for subsequent tests
with patch("transformers.models.hunyuan_v1_dense.modeling_hunyuan_v1_dense"):
from liger_kernel.transformers.model.hunyuan_v1 import lce_forward as hunyuan_v1_dense_lce_forward
# Instantiate a dummy model
config = transformers.models.hunyuan_v1_dense.configuration_hunyuan_v1_dense.HunYuanDenseV1Config(
torch_dtype=torch.bfloat16,
rms_norm_eps=1e-5,
hidden_size=32,
intermediate_size=64,
hidden_act="silu",
num_hidden_layers=2,
head_dim=1,
)
dummy_model_instance = AutoModelForCausalLM.from_config(config)
# Check that model instance variables are not yet patched with Liger modules
assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(hunyuan_v1_dense_lce_forward)
assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerSwiGLUMLP.forward)
assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
# Test applying kernels to the model instance
_apply_liger_kernel_to_instance(model=dummy_model_instance)
# Check that the model's instance variables were correctly patched with Liger modules
assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(hunyuan_v1_dense_lce_forward)
assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)
for layer in dummy_model_instance.model.layers:
assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerSwiGLUMLP.forward)
assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
try:
print(dummy_model_instance)
except Exception as e:
pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
+216
View File
@@ -0,0 +1,216 @@
import pytest
import torch
from test.utils import supports_bfloat16
from transformers.models.llama.configuration_llama import LlamaConfig
from liger_kernel.transformers.geglu import LigerGEGLUMLP
from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
from liger_kernel.transformers.tiled_mlp import LigerTiledGEGLUMLP
from liger_kernel.transformers.tiled_mlp import LigerTiledSwiGLUMLP
from liger_kernel.utils import infer_device
device = infer_device()
@pytest.mark.parametrize(
"bsz, seq_len, hidden_size, intermediate_size",
[
(1, 1024, 128, 256), # num_shards=8 if auto
(2, 1024, 64, 256), # num_shards=16 if auto
# weird shapes
(4, 127, 128, 256),
],
)
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
# atol is for small values: they have more difference, so set atol higher
# rtol is for larger values: they are very close, so set rtol lower
(torch.float32, 1e-0, 2e-6),
pytest.param(
torch.bfloat16,
1e-0,
1e-0,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
],
)
@pytest.mark.parametrize("num_shards", [None, 2, 4])
@pytest.mark.parametrize("check_2d", [True, False])
def test_tiled_geglu_correctness(bsz, seq_len, hidden_size, intermediate_size, dtype, atol, rtol, num_shards, check_2d):
"""Test that TiledGEGLUMLP produces similar results as regular GEGLUMLP."""
config = LlamaConfig(
hidden_size=hidden_size,
intermediate_size=intermediate_size,
hidden_act="gelu_pytorch_tanh",
)
# scale input so that the numerical errors are accumulated less
_input = torch.randn(bsz, seq_len, hidden_size, device=device, dtype=dtype) * 0.1
x1 = _input.detach().clone().requires_grad_(True)
x2 = _input.detach().clone().requires_grad_(True)
# Convert to 2D input for MoE experts testing
if check_2d:
x1 = x1.view(-1, hidden_size)
x2 = x2.view(-1, hidden_size)
# Initialize weights
G = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype)
U = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype)
D = torch.randn(hidden_size, intermediate_size, device=device, dtype=dtype)
# Regular GEGLU MLP
regular_mlp = LigerGEGLUMLP(config=config).to(device).to(dtype)
regular_mlp.gate_proj.weight.data = G
regular_mlp.up_proj.weight.data = U
regular_mlp.down_proj.weight.data = D
# Tiled GEGLU MLP
tiled_mlp = LigerTiledGEGLUMLP(config=config, num_shards=num_shards).to(device).to(dtype)
tiled_mlp.gate_proj.weight.data = G
tiled_mlp.up_proj.weight.data = U
tiled_mlp.down_proj.weight.data = D
# Forward pass
y1 = regular_mlp(x1)
y2 = tiled_mlp(x2)
torch.testing.assert_close(y1, y2, atol=atol, rtol=rtol, msg="Forward outputs don't match")
# Backward pass
dy = torch.randn_like(y1)
y1.backward(dy.clone(), retain_graph=True)
y2.backward(dy.clone(), retain_graph=True)
# Check gradients
torch.testing.assert_close(
regular_mlp.gate_proj.weight.grad,
tiled_mlp.gate_proj.weight.grad,
atol=atol,
rtol=rtol,
msg="gate_proj weight gradients don't match",
)
torch.testing.assert_close(
regular_mlp.up_proj.weight.grad,
tiled_mlp.up_proj.weight.grad,
atol=atol,
rtol=rtol,
msg="up_proj weight gradients don't match",
)
torch.testing.assert_close(
regular_mlp.down_proj.weight.grad,
tiled_mlp.down_proj.weight.grad,
atol=atol,
rtol=rtol,
msg="down_proj weight gradients don't match",
)
torch.testing.assert_close(x1.grad, x2.grad, atol=atol, rtol=rtol, msg="Input gradients don't match")
@pytest.mark.parametrize(
"bsz, seq_len, hidden_size, intermediate_size",
[
(2, 512, 512, 1024),
(1, 1024, 256, 512),
# weird shapes
(4, 127, 128, 256),
],
)
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
# atol is for small values: they have more difference, so set atol higher
# rtol is for larger values: they are very close, so set rtol lower
(torch.float32, 1e-0, 2e-6),
pytest.param(
torch.bfloat16,
1e-0,
1e-0,
marks=pytest.mark.skip(reason="bfloat16 tests disabled due to numerical instability"),
),
],
)
@pytest.mark.parametrize("num_shards", [None, 2, 4])
@pytest.mark.parametrize("check_2d", [True, False])
def test_tiled_swiglu_correctness(
bsz, seq_len, hidden_size, intermediate_size, dtype, atol, rtol, num_shards, check_2d
):
"""Test that TiledSwiGLUMLP produces similar results as regular SwiGLUMLP."""
config = LlamaConfig(
hidden_size=hidden_size,
intermediate_size=intermediate_size,
hidden_act="silu",
)
# scale input so that the numerical errors are accumulated less
_input = torch.randn(bsz, seq_len, hidden_size, device=device, dtype=dtype) * 0.1
x1 = _input.detach().clone().requires_grad_(True)
x2 = _input.detach().clone().requires_grad_(True)
# Convert to 2D input for MoE experts testing
if check_2d:
x1 = x1.view(-1, hidden_size)
x2 = x2.view(-1, hidden_size)
# Initialize weights
G = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype)
U = torch.randn(intermediate_size, hidden_size, device=device, dtype=dtype)
D = torch.randn(hidden_size, intermediate_size, device=device, dtype=dtype)
# Regular SwiGLU MLP
regular_mlp = LigerSwiGLUMLP(config=config).to(device).to(dtype)
regular_mlp.gate_proj.weight.data = G
regular_mlp.up_proj.weight.data = U
regular_mlp.down_proj.weight.data = D
# Tiled SwiGLU MLP
tiled_mlp = LigerTiledSwiGLUMLP(config=config, num_shards=num_shards).to(device).to(dtype)
tiled_mlp.gate_proj.weight.data = G
tiled_mlp.up_proj.weight.data = U
tiled_mlp.down_proj.weight.data = D
# Forward pass
y1 = regular_mlp(x1)
y2 = tiled_mlp(x2)
torch.testing.assert_close(y1, y2, atol=atol, rtol=rtol, msg="Forward outputs don't match")
# Backward pass
dy = torch.randn_like(y1)
y1.backward(dy.clone(), retain_graph=True)
y2.backward(dy.clone(), retain_graph=True)
# Check gradients
torch.testing.assert_close(
regular_mlp.gate_proj.weight.grad,
tiled_mlp.gate_proj.weight.grad,
atol=atol,
rtol=rtol,
msg="gate_proj weight gradients don't match",
)
torch.testing.assert_close(
regular_mlp.up_proj.weight.grad,
tiled_mlp.up_proj.weight.grad,
atol=atol,
rtol=rtol,
msg="up_proj weight gradients don't match",
)
torch.testing.assert_close(
regular_mlp.down_proj.weight.grad,
tiled_mlp.down_proj.weight.grad,
atol=atol,
rtol=rtol,
msg="down_proj weight gradients don't match",
)
torch.testing.assert_close(x1.grad, x2.grad, atol=atol, rtol=rtol, msg="Input gradients don't match")
+36
View File
@@ -550,6 +550,18 @@ def revert_liger_kernel_to_olmo2(model_config: MiniModelConfig):
print("Liger kernel patches have been reverted.")
def revert_liger_kernel_to_olmo3(model_config: MiniModelConfig):
"""
Revert all Liger kernel patches applied to Olmo3.
"""
from transformers.models.olmo3 import modeling_olmo3
importlib.reload(modeling_olmo3)
model_config.model_class = modeling_olmo3.Olmo3ForCausalLM
print("Liger kernel patches have been reverted.")
def revert_liger_kernel_to_glm4(model_config: MiniModelConfig):
"""
Revert all Liger kernel patches applied to Glm4.
@@ -661,6 +673,30 @@ def revert_liger_kernel_to_qwen3_next(model_config: MiniModelConfig):
print("Liger kernel patches have been reverted.")
def revert_liger_kernel_to_hunyuan_v1(model_config: MiniModelConfig):
"""
Revert all Liger kernel patches applied to Hunyuanv1.
"""
from transformers.models.hunyuan_v1_dense import modeling_hunyuan_v1_dense
importlib.reload(modeling_hunyuan_v1_dense)
model_config.model_class = modeling_hunyuan_v1_dense.HunYuanDenseV1ForCausalLM
print("Liger kernel patches have been reverted.")
def revert_liger_kernel_to_hunyuan_v1_moe(model_config: MiniModelConfig):
"""
Revert all Liger kernel patches applied to Hunyuanv1 MoE.
"""
from transformers.models.hunyuan_v1_moe import modeling_hunyuan_v1_moe
importlib.reload(modeling_hunyuan_v1_moe)
model_config.model_class = modeling_hunyuan_v1_moe.HunYuanMoEV1ForCausalLM
print("Liger kernel patches have been reverted.")
class HFAlignmentLoss:
def __init__(
self,