From 6bf2e5ad3ed717385afbdf7149528ce3f5a31e58 Mon Sep 17 00:00:00 2001 From: Buffett Liu Date: Sun, 26 Jul 2026 16:57:53 -0700 Subject: [PATCH 1/2] [Megatron] Add SwiGLU integration --- benchmark/data/all_benchmark_data.csv | 96 +++ .../scripts/benchmark_megatron_swiglu.py | 186 +++++ examples/megatron/README.md | 38 +- examples/megatron/run_mode1_monkey_patch.py | 71 +- examples/megatron/run_mode2_hand_spec.py | 99 ++- src/liger_kernel/megatron/__init__.py | 10 +- src/liger_kernel/megatron/monkey_patch.py | 139 +++- src/liger_kernel/megatron/swiglu.py | 162 +++++ src/liger_kernel/ops/__init__.py | 1 + src/liger_kernel/ops/swiglu.py | 145 ++++ test/megatron/test_monkey_patch.py | 464 ++++++++++++- test/megatron/test_swiglu.py | 642 ++++++++++++++++++ 12 files changed, 1997 insertions(+), 56 deletions(-) create mode 100644 benchmark/scripts/benchmark_megatron_swiglu.py create mode 100644 src/liger_kernel/megatron/swiglu.py create mode 100644 test/megatron/test_swiglu.py diff --git a/benchmark/data/all_benchmark_data.csv b/benchmark/data/all_benchmark_data.csv index 84e73583c..196e2d97b 100644 --- a/benchmark/data/all_benchmark_data.csv +++ b/benchmark/data/all_benchmark_data.csv @@ -2351,3 +2351,99 @@ megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,16384,1024.1 megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,32768,2048.16552734375,2048.16552734375,2048.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0 megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,65536,4096.16552734375,4096.16552734375,4096.16552734375,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0 megatron_cross_entropy,megatron-unfused,full,memory,MB,V,vocab size,131072,8192.166015625,8192.166015625,8192.166015625,"{""S"": 2048, ""B"": 4, ""TP"": 1}",NVIDIA H100 80GB HBM3,2026-06-15 22:13:03,0.8.0 +megatron_swiglu,liger,forward,speed,ms,ffn_local,per-rank FFN hidden size,1024,0.02332800067961216,0.02300800010561943,0.023903999477624893,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:43,0.8.1 +megatron_swiglu,liger,forward,speed,ms,ffn_local,per-rank FFN hidden size,2048,0.039583999663591385,0.03923200070858002,0.04053119942545891,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:43,0.8.1 +megatron_swiglu,liger,forward,speed,ms,ffn_local,per-rank FFN hidden size,4096,0.07251200079917908,0.07161600142717361,0.0728704035282135,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:43,0.8.1 +megatron_swiglu,liger,forward,speed,ms,ffn_local,per-rank FFN hidden size,8192,0.13708800077438354,0.1366720050573349,0.13788799941539764,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:43,0.8.1 +megatron_swiglu,liger,forward,speed,ms,ffn_local,per-rank FFN hidden size,16384,0.2670080065727234,0.26611199975013733,0.2675136089324951,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:43,0.8.1 +megatron_swiglu,liger,forward,speed,ms,ffn_local,per-rank FFN hidden size,32768,0.5275039970874786,0.526918375492096,0.5280320048332214,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:43,0.8.1 +megatron_swiglu,liger_in_place,forward,speed,ms,ffn_local,per-rank FFN hidden size,1024,0.023135999217629433,0.022816000506281853,0.02367359958589077,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:44,0.8.1 +megatron_swiglu,liger_in_place,forward,speed,ms,ffn_local,per-rank FFN hidden size,2048,0.03964800015091896,0.03929600119590759,0.04057599976658821,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:44,0.8.1 +megatron_swiglu,liger_in_place,forward,speed,ms,ffn_local,per-rank FFN hidden size,4096,0.07264000177383423,0.07180800288915634,0.07302399724721909,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:44,0.8.1 +megatron_swiglu,liger_in_place,forward,speed,ms,ffn_local,per-rank FFN hidden size,8192,0.13699199259281158,0.1366720050573349,0.13788799941539764,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:44,0.8.1 +megatron_swiglu,liger_in_place,forward,speed,ms,ffn_local,per-rank FFN hidden size,16384,0.2667680084705353,0.26579201221466064,0.2673279941082001,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:44,0.8.1 +megatron_swiglu,liger_in_place,forward,speed,ms,ffn_local,per-rank FFN hidden size,32768,0.5273440182209015,0.5268160104751587,0.5279743909835816,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:44,0.8.1 +megatron_swiglu,torch,forward,speed,ms,ffn_local,per-rank FFN hidden size,1024,0.05782400071620941,0.05718399956822395,0.05904639884829521,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:45,0.8.1 +megatron_swiglu,torch,forward,speed,ms,ffn_local,per-rank FFN hidden size,2048,0.10529599711298943,0.10502400249242783,0.10695680230855942,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:45,0.8.1 +megatron_swiglu,torch,forward,speed,ms,ffn_local,per-rank FFN hidden size,4096,0.19763199985027313,0.19737599790096283,0.19925120174884797,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:45,0.8.1 +megatron_swiglu,torch,forward,speed,ms,ffn_local,per-rank FFN hidden size,8192,0.38304001092910767,0.38252800703048706,0.3841023981571198,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:45,0.8.1 +megatron_swiglu,torch,forward,speed,ms,ffn_local,per-rank FFN hidden size,16384,0.7503359913825989,0.7496960163116455,0.7515007734298705,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:45,0.8.1 +megatron_swiglu,torch,forward,speed,ms,ffn_local,per-rank FFN hidden size,32768,1.4952960014343262,1.494860792160034,1.4964608430862425,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:45,0.8.1 +megatron_swiglu,megatron,forward,speed,ms,ffn_local,per-rank FFN hidden size,1024,0.02319999970495701,0.02300800010561943,0.023744000121951103,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:48,0.8.1 +megatron_swiglu,megatron,forward,speed,ms,ffn_local,per-rank FFN hidden size,2048,0.060416001826524734,0.060095999389886856,0.060736000537872314,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:48,0.8.1 +megatron_swiglu,megatron,forward,speed,ms,ffn_local,per-rank FFN hidden size,4096,0.11369600147008896,0.11343999952077866,0.11401599645614624,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:48,0.8.1 +megatron_swiglu,megatron,forward,speed,ms,ffn_local,per-rank FFN hidden size,8192,0.22015999257564545,0.21992959976196289,0.22044800221920013,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:48,0.8.1 +megatron_swiglu,megatron,forward,speed,ms,ffn_local,per-rank FFN hidden size,16384,0.4328640103340149,0.4325439929962158,0.4331200122833252,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:48,0.8.1 +megatron_swiglu,megatron,forward,speed,ms,ffn_local,per-rank FFN hidden size,32768,0.8584960103034973,0.8581759929656982,0.8588160276412964,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:48,0.8.1 +megatron_swiglu,liger,backward,speed,ms,ffn_local,per-rank FFN hidden size,1024,0.10713600367307663,0.09860479831695557,0.11654400080442429,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:50,0.8.1 +megatron_swiglu,liger,backward,speed,ms,ffn_local,per-rank FFN hidden size,2048,0.14075200259685516,0.13971200585365295,0.14217600226402283,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:50,0.8.1 +megatron_swiglu,liger,backward,speed,ms,ffn_local,per-rank FFN hidden size,4096,0.2667520046234131,0.2662400007247925,0.2673279941082001,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:50,0.8.1 +megatron_swiglu,liger,backward,speed,ms,ffn_local,per-rank FFN hidden size,8192,0.5025280117988586,0.5018048167228699,0.503219199180603,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:50,0.8.1 +megatron_swiglu,liger,backward,speed,ms,ffn_local,per-rank FFN hidden size,16384,0.9709759950637817,0.9701248049736022,0.9725504159927367,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:50,0.8.1 +megatron_swiglu,liger,backward,speed,ms,ffn_local,per-rank FFN hidden size,32768,1.9115840196609497,1.9104703903198244,1.9125759601593018,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:50,0.8.1 +megatron_swiglu,liger_in_place,backward,speed,ms,ffn_local,per-rank FFN hidden size,1024,0.09843199700117111,0.09151999801397324,0.10835839807987213,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:50,0.8.1 +megatron_swiglu,liger_in_place,backward,speed,ms,ffn_local,per-rank FFN hidden size,2048,0.14124800264835358,0.14019200205802917,0.1423552006483078,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:50,0.8.1 +megatron_swiglu,liger_in_place,backward,speed,ms,ffn_local,per-rank FFN hidden size,4096,0.26499199867248535,0.26371839046478274,0.2657279968261719,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:50,0.8.1 +megatron_swiglu,liger_in_place,backward,speed,ms,ffn_local,per-rank FFN hidden size,8192,0.5005760192871094,0.49983999133110046,0.5016000270843506,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:50,0.8.1 +megatron_swiglu,liger_in_place,backward,speed,ms,ffn_local,per-rank FFN hidden size,16384,0.9724159836769104,0.9714111804962158,0.9733887791633606,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:50,0.8.1 +megatron_swiglu,liger_in_place,backward,speed,ms,ffn_local,per-rank FFN hidden size,32768,1.9130080342292786,1.91212797164917,1.9141120195388794,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:50,0.8.1 +megatron_swiglu,torch,backward,speed,ms,ffn_local,per-rank FFN hidden size,1024,0.17257599532604218,0.1706559956073761,0.17401599884033203,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:51,0.8.1 +megatron_swiglu,torch,backward,speed,ms,ffn_local,per-rank FFN hidden size,2048,0.317984014749527,0.3160511970520019,0.3192128002643585,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:51,0.8.1 +megatron_swiglu,torch,backward,speed,ms,ffn_local,per-rank FFN hidden size,4096,0.6032639741897583,0.6014080047607422,0.6043840050697327,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:51,0.8.1 +megatron_swiglu,torch,backward,speed,ms,ffn_local,per-rank FFN hidden size,8192,1.1695040464401245,1.167942428588867,1.1714047908782959,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:51,0.8.1 +megatron_swiglu,torch,backward,speed,ms,ffn_local,per-rank FFN hidden size,16384,2.2973440885543823,2.2952704429626465,2.298630380630493,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:51,0.8.1 +megatron_swiglu,torch,backward,speed,ms,ffn_local,per-rank FFN hidden size,32768,4.572864055633545,4.571807861328125,4.574592113494873,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:51,0.8.1 +megatron_swiglu,megatron,backward,speed,ms,ffn_local,per-rank FFN hidden size,1024,0.14691199362277985,0.1360447973012924,0.15591679811477663,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:55,0.8.1 +megatron_swiglu,megatron,backward,speed,ms,ffn_local,per-rank FFN hidden size,2048,0.33847999572753906,0.3375552117824554,0.3397376000881195,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:55,0.8.1 +megatron_swiglu,megatron,backward,speed,ms,ffn_local,per-rank FFN hidden size,4096,0.6568639874458313,0.6558079719543457,0.6579840183258057,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:55,0.8.1 +megatron_swiglu,megatron,backward,speed,ms,ffn_local,per-rank FFN hidden size,8192,1.2886880040168762,1.287782382965088,1.289689588546753,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:55,0.8.1 +megatron_swiglu,megatron,backward,speed,ms,ffn_local,per-rank FFN hidden size,16384,2.554271936416626,2.5534400939941406,2.5555647850036625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:55,0.8.1 +megatron_swiglu,megatron,backward,speed,ms,ffn_local,per-rank FFN hidden size,32768,5.0830078125,5.082399845123291,5.084876918792725,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:55,0.8.1 +megatron_swiglu,liger,full,speed,ms,ffn_local,per-rank FFN hidden size,1024,0.12219199910759926,0.11427199840545654,0.13155840039253236,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:55,0.8.1 +megatron_swiglu,liger,full,speed,ms,ffn_local,per-rank FFN hidden size,2048,0.14047999680042267,0.1393599957227707,0.1414719969034195,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:55,0.8.1 +megatron_swiglu,liger,full,speed,ms,ffn_local,per-rank FFN hidden size,4096,0.26678401231765747,0.2662464022636414,0.26736000180244446,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:55,0.8.1 +megatron_swiglu,liger,full,speed,ms,ffn_local,per-rank FFN hidden size,8192,0.5025920271873474,0.5019519925117493,0.503167986869812,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:55,0.8.1 +megatron_swiglu,liger,full,speed,ms,ffn_local,per-rank FFN hidden size,16384,0.9714719951152802,0.9705024003982544,0.9724992036819458,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:55,0.8.1 +megatron_swiglu,liger,full,speed,ms,ffn_local,per-rank FFN hidden size,32768,1.9120799899101257,1.9109375953674317,1.9139583826065063,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:55,0.8.1 +megatron_swiglu,liger_in_place,full,speed,ms,ffn_local,per-rank FFN hidden size,1024,0.12014399841427803,0.11004160344600677,0.1353279948234558,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:56,0.8.1 +megatron_swiglu,liger_in_place,full,speed,ms,ffn_local,per-rank FFN hidden size,2048,0.14028799533843994,0.13913600146770477,0.14156800508499146,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:56,0.8.1 +megatron_swiglu,liger_in_place,full,speed,ms,ffn_local,per-rank FFN hidden size,4096,0.26633599400520325,0.2659519910812378,0.26688000559806824,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:56,0.8.1 +megatron_swiglu,liger_in_place,full,speed,ms,ffn_local,per-rank FFN hidden size,8192,0.5021439790725708,0.501695990562439,0.5027199983596802,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:56,0.8.1 +megatron_swiglu,liger_in_place,full,speed,ms,ffn_local,per-rank FFN hidden size,16384,0.9710400104522705,0.9700352072715759,0.9721791863441467,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:56,0.8.1 +megatron_swiglu,liger_in_place,full,speed,ms,ffn_local,per-rank FFN hidden size,32768,1.9120000004768372,1.9105535984039308,1.9125887632369996,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:56,0.8.1 +megatron_swiglu,torch,full,speed,ms,ffn_local,per-rank FFN hidden size,1024,0.17270399630069733,0.17081600427627563,0.1740799993276596,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:57,0.8.1 +megatron_swiglu,torch,full,speed,ms,ffn_local,per-rank FFN hidden size,2048,0.31836798787117004,0.3162751913070679,0.31934080123901365,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:57,0.8.1 +megatron_swiglu,torch,full,speed,ms,ffn_local,per-rank FFN hidden size,4096,0.6027200222015381,0.6009856224060058,0.6041920185089111,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:57,0.8.1 +megatron_swiglu,torch,full,speed,ms,ffn_local,per-rank FFN hidden size,8192,1.1698559522628784,1.1680512189865113,1.1706367969512939,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:57,0.8.1 +megatron_swiglu,torch,full,speed,ms,ffn_local,per-rank FFN hidden size,16384,2.2963199615478516,2.294924736022949,2.2984512329101565,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:57,0.8.1 +megatron_swiglu,torch,full,speed,ms,ffn_local,per-rank FFN hidden size,32768,4.5731520652771,4.5726399421691895,4.574336051940918,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:57,0.8.1 +megatron_swiglu,megatron,full,speed,ms,ffn_local,per-rank FFN hidden size,1024,0.1780479997396469,0.17710720300674437,0.18151679933071138,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,megatron,full,speed,ms,ffn_local,per-rank FFN hidden size,2048,0.3383679986000061,0.33756161332130435,0.33961600065231323,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,megatron,full,speed,ms,ffn_local,per-rank FFN hidden size,4096,0.6570240259170532,0.6558911919593811,0.6581119894981384,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,megatron,full,speed,ms,ffn_local,per-rank FFN hidden size,8192,1.2888000011444092,1.2880703926086425,1.2902079820632935,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,megatron,full,speed,ms,ffn_local,per-rank FFN hidden size,16384,2.554080009460449,2.553056001663208,2.555039882659912,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,megatron,full,speed,ms,ffn_local,per-rank FFN hidden size,32768,5.084127902984619,5.083379173278809,5.085017681121826,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,liger,full,memory,MB,ffn_local,per-rank FFN hidden size,1024,96.0009765625,96.0009765625,96.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,liger,full,memory,MB,ffn_local,per-rank FFN hidden size,2048,192.0009765625,192.0009765625,192.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,liger,full,memory,MB,ffn_local,per-rank FFN hidden size,4096,384.0009765625,384.0009765625,384.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,liger,full,memory,MB,ffn_local,per-rank FFN hidden size,8192,768.0009765625,768.0009765625,768.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,liger,full,memory,MB,ffn_local,per-rank FFN hidden size,16384,1536.0009765625,1536.0009765625,1536.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,liger,full,memory,MB,ffn_local,per-rank FFN hidden size,32768,3072.0009765625,3072.0009765625,3072.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,liger_in_place,full,memory,MB,ffn_local,per-rank FFN hidden size,1024,64.0009765625,64.0009765625,64.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,liger_in_place,full,memory,MB,ffn_local,per-rank FFN hidden size,2048,128.0009765625,128.0009765625,128.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,liger_in_place,full,memory,MB,ffn_local,per-rank FFN hidden size,4096,256.0009765625,256.0009765625,256.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,liger_in_place,full,memory,MB,ffn_local,per-rank FFN hidden size,8192,512.0009765625,512.0009765625,512.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,liger_in_place,full,memory,MB,ffn_local,per-rank FFN hidden size,16384,1024.0009765625,1024.0009765625,1024.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,liger_in_place,full,memory,MB,ffn_local,per-rank FFN hidden size,32768,2048.0009765625,2048.0009765625,2048.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,torch,full,memory,MB,ffn_local,per-rank FFN hidden size,1024,112.0009765625,112.0009765625,112.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,torch,full,memory,MB,ffn_local,per-rank FFN hidden size,2048,224.0009765625,224.0009765625,224.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,torch,full,memory,MB,ffn_local,per-rank FFN hidden size,4096,448.0009765625,448.0009765625,448.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,torch,full,memory,MB,ffn_local,per-rank FFN hidden size,8192,896.0009765625,896.0009765625,896.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,torch,full,memory,MB,ffn_local,per-rank FFN hidden size,16384,1792.0009765625,1792.0009765625,1792.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,torch,full,memory,MB,ffn_local,per-rank FFN hidden size,32768,3584.0009765625,3584.0009765625,3584.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,megatron,full,memory,MB,ffn_local,per-rank FFN hidden size,1024,80.0009765625,80.0009765625,80.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,megatron,full,memory,MB,ffn_local,per-rank FFN hidden size,2048,160.0009765625,160.0009765625,160.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,megatron,full,memory,MB,ffn_local,per-rank FFN hidden size,4096,320.0009765625,320.0009765625,320.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,megatron,full,memory,MB,ffn_local,per-rank FFN hidden size,8192,640.0009765625,640.0009765625,640.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,megatron,full,memory,MB,ffn_local,per-rank FFN hidden size,16384,1280.0009765625,1280.0009765625,1280.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 +megatron_swiglu,megatron,full,memory,MB,ffn_local,per-rank FFN hidden size,32768,2560.0009765625,2560.0009765625,2560.0009765625,"{""S"": 2048, ""B"": 4}",NVIDIA H100 80GB HBM3,2026-07-26 00:09:58,0.8.1 diff --git a/benchmark/scripts/benchmark_megatron_swiglu.py b/benchmark/scripts/benchmark_megatron_swiglu.py new file mode 100644 index 000000000..48e037683 --- /dev/null +++ b/benchmark/scripts/benchmark_megatron_swiglu.py @@ -0,0 +1,186 @@ +"""Benchmark Liger's Megatron-LM SwiGLU wrapper. + +Compares four providers on the fused gate-up activation call shape +``[seq, batch, 2 * ffn_local]`` -- exactly the tensor Megatron's ``linear_fc1`` hands to +``bias_swiglu_impl``: + + - **liger**: ``LigerMegatronSwiGLU`` -- Liger's fused gate-up Triton kernel in the + Megatron-shaped wrapper. This is the default configuration. + - **liger_in_place**: the same, with ``in_place=True``, which writes the backward + gradient into the fc1 output buffer rather than allocating a new one. Same speed, + one fewer activation-sized allocation; opt-in because it destroys that buffer. + - **megatron**: Megatron's ``bias_swiglu_impl`` (``bias_activation_fusion=True``), a + chain of ``@jit_fuser`` TorchScript helpers. This is the symbol Liger displaces. + - **torch**: eager ``F.silu(y_1) * y_2`` over ``torch.chunk(y, 2, -1)`` -- the + unfused reference, and structurally what Megatron runs when + ``bias_activation_fusion=False`` (there it's a closure inside ``MLP.forward``). + +Why there is no ``--tp-size`` flag (unlike the Megatron CE benchmark): SwiGLU is +elementwise and token-local. It issues **no collectives**, and tensor parallelism affects +it only by shrinking the per-rank column count to ``ffn_hidden_size / tp``. Sweeping +``ffn_local`` on a single GPU therefore already covers every TP configuration -- TP=8 at +``ffn=28672`` is the same kernel work as the ``ffn_local=3584`` point on this curve. Peak +memory is likewise per-rank and scales as 1/TP. + +The x-axis spans the Blackwell tiling threshold: ``liger_kernel.ops.swiglu`` switches to a +column-tiled 2D grid when ``next_pow2(n_cols) >= 16384`` on Blackwell, so the two largest +points exercise that path on B200 and the one-row path everywhere else. + +Requires a Liger-supported accelerator (CUDA / ROCm). With megatron-core not installed the +``megatron`` provider is silently dropped and the run proceeds with ``liger`` + ``torch``. + +Output goes to the shared ``benchmark/data/all_benchmark_data.csv`` -- rows are tagged with +``kernel_name="megatron_swiglu"`` and the standard visualizer renders them via: + + python benchmark/benchmarks_visualizer.py \\ + --kernel-name megatron_swiglu --metric-name speed + python benchmark/benchmarks_visualizer.py \\ + --kernel-name megatron_swiglu --metric-name memory +""" + +import torch +import torch.nn.functional as F +import triton + +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.megatron import LigerMegatronSwiGLU +from liger_kernel.utils import infer_device + +device = infer_device() + +try: + from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl + + _MEGATRON_AVAILABLE = True +except ImportError: + bias_swiglu_impl = None + _MEGATRON_AVAILABLE = False + + +def _torch_swiglu(y): + """Eager reference — identical math to Megatron's ``swiglu``, minus the JIT fusion.""" + y_1, y_2 = torch.chunk(y, 2, -1) + return F.silu(y_1) * y_2 + + +def _make_fwd(provider: str): + if provider == "liger": + module = LigerMegatronSwiGLU() + return lambda y: module(y, None, False, False) + if provider == "liger_in_place": + module = LigerMegatronSwiGLU(in_place=True) + return lambda y: module(y, None, False, False) + if provider == "torch": + return _torch_swiglu + if provider == "megatron": + if not _MEGATRON_AVAILABLE: + raise RuntimeError("megatron-core not installed; cannot benchmark 'megatron' provider") + return lambda y: bias_swiglu_impl(y, None, False, False) + raise ValueError(f"unknown provider: {provider!r}") + + +def _make_input(s: int, b: int, ffn_local: int, requires_grad: bool = True) -> torch.Tensor: + # 2 * ffn_local: Megatron's linear_fc1 emits gate and up concatenated on the last dim. + return torch.randn(s, b, 2 * ffn_local, device=device, dtype=torch.bfloat16, requires_grad=requires_grad) + + +def bench_speed_megatron_swiglu(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + ffn_local = input.x + provider = input.kernel_provider + mode = input.kernel_operation_mode + s = input.extra_benchmark_config["S"] + b = input.extra_benchmark_config["B"] + + fwd_fn = _make_fwd(provider) + x = _make_input(s, b, ffn_local) + + def fwd(): + return fwd_fn(x) + + if mode == "forward": + ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) + elif mode == "backward": + # Rerun fwd each iteration: Liger's in-place backward consumes the saved buffers, + # so a retained graph would corrupt on the second pass. Subtract the "forward" + # row to get backward-only timing. + def _fwd_bwd(): + if x.grad is not None: + x.grad = None + out = fwd() + out.sum().backward() + + ms_50, ms_20, ms_80 = triton.testing.do_bench(_fwd_bwd, rep=100, quantiles=QUANTILES) + elif mode == "full": + + def full(): + if x.grad is not None: + x.grad = None + y = fwd() + y.sum().backward() + + ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES) + else: + raise ValueError(f"unknown mode: {mode!r}") + + return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80) + + +def bench_memory_megatron_swiglu(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + ffn_local = input.x + provider = input.kernel_provider + s = input.extra_benchmark_config["S"] + b = input.extra_benchmark_config["B"] + + fwd_fn = _make_fwd(provider) + x = _make_input(s, b, ffn_local) + + def full(): + if x.grad is not None: + x.grad = None + y = fwd_fn(x) + y.sum().backward() + + 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() + + providers = ["liger", "liger_in_place", "torch"] + if _MEGATRON_AVAILABLE: + providers.append("megatron") + + common_configs = { + "kernel_name": "megatron_swiglu", + "x_name": "ffn_local", + "x_label": "per-rank FFN hidden size", + # 1024 → 32768. Llama-7B is 11008 and Llama-70B is 28672, so this brackets + # production sizes; the top two points cross the Blackwell tiling threshold. + "x_values": [2**i for i in range(10, 16)], + "kernel_providers": providers, + # Megatron's standard training shape, matching the megatron CE benchmark. + "extra_benchmark_configs": [{"S": 2048, "B": 4}], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_megatron_swiglu, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_megatron_swiglu, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/examples/megatron/README.md b/examples/megatron/README.md index 0f0d247e8..126dc7888 100644 --- a/examples/megatron/README.md +++ b/examples/megatron/README.md @@ -2,8 +2,9 @@ Two self-contained scripts demonstrating the integration modes shipped by `liger_kernel.megatron`. Both train a tiny GPT model with mock data on -2 × GPU (TP=2, PP=1) for 5 iterations and print the resolved norm classes -so you can see which slots picked up Liger. +2 × GPU (TP=1, PP=1, DP=2) for 5 iterations and print the resolved norm +classes, CE bindings and SwiGLU bindings so you can see which slots picked +up Liger. ## Prerequisites @@ -15,8 +16,15 @@ so you can see which slots picked up Liger. ## Mode 1 — `apply_liger_kernel_to_megatron()` (monkey-patch) One-line opt-in. Patches `LocalSpecProvider.layer_norm` and -`transformer_block.LayerNormImpl` so every RMSNorm slot becomes Liger -without changing the spec the user constructs. +`transformer_block.LayerNormImpl` so every RMSNorm slot becomes Liger, the +two cross-entropy entry points, and `bias_swiglu_impl` — all without +changing the spec the user constructs. + +Note the SwiGLU dispatch flags on `TransformerConfig` +(`gated_linear_unit=True`, `activation_func=F.silu`, +`bias_activation_fusion=True`). `MLP.forward` only routes through +`bias_swiglu_impl` when all three hold, so without them the patch is +applied but never reached — no speedup and no error. ```bash torchrun --nproc_per_node=2 \ @@ -27,9 +35,16 @@ torchrun --nproc_per_node=2 \ ## Mode 2 — hand-assembled `TransformerBlockSubmodules` Slot-level control. Explicitly places `LigerMegatronRMSNorm` into each -norm slot, including the block-level `final_layernorm`. Useful when you -want to mix Liger with other backends (e.g. TransformerEngine) on a -per-slot basis. +norm slot, including the block-level `final_layernorm`, and a +`LigerMegatronSwiGLU`-backed `MLP` subclass into the `mlp` slot. Useful +when you want to mix Liger with other backends (e.g. TransformerEngine) on +a per-slot basis. + +SwiGLU has no spec slot of its own — `MLPSubmodules.activation_func` is +only read when `config.use_te_activation_func` is set, so it is the +TransformerEngine hook rather than a general one. The MLP *class* is a +slot, so this script subclasses `MLP`, the same way it subclasses +`GPTModel` for cross-entropy. ```bash torchrun --nproc_per_node=2 \ @@ -41,13 +56,20 @@ torchrun --nproc_per_node=2 \ For both scripts: -- 5 lines of `[modeN] iter loss=` with the loss decreasing. +- 5 lines of `[modeN] iter loss=`. (Five iterations on a + 12-hidden-size model with random mock data is far too short to show a + trend — the value just hovers.) - A printed module tree with `LigerMegatronRMSNorm` in **5 of 5** norm slots: four per-layer (`input_layernorm`, `pre_mlp_layernorm` × 2 layers) and one block-level (`final_layernorm`). - `Successfully loaded the model` after the distributed checkpoint round-trip. +Mode 1 additionally prints `=== Resolved SwiGLU symbols ===` with all three +bindings tagged `[Liger]` — the defining module plus the two consumers that +import the symbol by name. Mode 2 prints `=== Resolved SwiGLU MLPs ===` +listing both `_LigerSwiGLUMLP` layers. + If your environment doesn't have Apex or TransformerEngine installed, you will see harmless warnings — Megatron falls back to the local backend, which is exactly where Liger plugs in. diff --git a/examples/megatron/run_mode1_monkey_patch.py b/examples/megatron/run_mode1_monkey_patch.py index 5e3f460df..64563d493 100644 --- a/examples/megatron/run_mode1_monkey_patch.py +++ b/examples/megatron/run_mode1_monkey_patch.py @@ -1,24 +1,35 @@ -"""Mode 1 — monkey-patch Megatron-Core to use Liger RMSNorm + cross-entropy. +"""Mode 1 — monkey-patch Megatron-Core to use Liger RMSNorm + cross-entropy + SwiGLU. Adapted from Megatron's ``examples/run_simple_mcore_train_loop.py``. The relevant additions (vs. that file) are: - 1. ``apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True)`` - called once at the top of ``model_provider()``. This patches: + 1. ``apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True, + swiglu=True)`` called once at the top of ``model_provider()``. This patches: - ``LocalSpecProvider.layer_norm`` (per-layer norm slots) - ``transformer_block.LayerNormImpl`` (block-level ``final_layernorm``) - ``fused_cross_entropy.fused_vocab_parallel_cross_entropy`` (the fused CE path) - ``tensor_parallel.cross_entropy.vocab_parallel_cross_entropy`` (the unfused CE path) + - ``fusions.fused_bias_swiglu.bias_swiglu_impl`` *and* the by-name + re-exports of it in ``transformer.mlp`` and + ``transformer.moe.shared_experts`` (the SwiGLU activation) 2. ``normalization="RMSNorm"`` added to ``TransformerConfig`` so the model actually has RMSNorm slots to patch (Megatron defaults to ``LayerNorm``). - 3. ``_print_norm_classes`` + ``_print_ce_symbols`` after model construction - — print the resolved class/function bindings so you can verify Liger - took over for every slot. + 3. The SwiGLU dispatch flags added to ``TransformerConfig`` + (``gated_linear_unit``, ``activation_func=F.silu``, + ``bias_activation_fusion``). ``MLP.forward`` only routes through + ``bias_swiglu_impl`` when all three hold, so without them the patch is + applied but never reached. ``add_bias_linear=False`` keeps ``bias`` at + ``None``, which is the configuration Liger accelerates; with a bias the + wrapper transparently falls back to Megatron's own kernel. + + 4. ``_print_norm_classes`` + ``_print_ce_symbols`` + ``_print_swiglu_symbols`` + after model construction — print the resolved class/function bindings so + you can verify Liger took over for every slot. Run with: torchrun --nproc_per_node=2 --master_addr=127.0.0.1 --master_port=29500 \\ @@ -34,6 +45,7 @@ from typing import Iterator import torch +import torch.nn.functional as F from megatron.core import dist_checkpointing from megatron.core import parallel_state @@ -74,7 +86,7 @@ def initialize_distributed(tp: int = 2, pp: int = 1) -> None: def model_provider() -> GPTModel: # ↓↓ Mode 1 — patch once, everything below picks up Liger ↓↓ - apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True) + apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True, swiglu=True) # ↑↑ ------------------------------------------------------ ↑↑ cfg = TransformerConfig( @@ -84,6 +96,10 @@ def model_provider() -> GPTModel: use_cpu_initialization=True, pipeline_dtype=torch.float32, normalization="RMSNorm", + gated_linear_unit=True, + activation_func=F.silu, + bias_activation_fusion=True, + add_bias_linear=False, ) return GPTModel( config=cfg, @@ -93,9 +109,23 @@ def model_provider() -> GPTModel: ) +def _ensure_dataset_helpers() -> None: + """Build Megatron's C++ dataset helpers only if not already compiled. + + ``compile_helpers()`` shells out to ``make``, which fails on pip-installed + megatron-core wheels (no Makefile, but the extension is already present). + """ + try: + from megatron.core.datasets import helpers_cpp # noqa: F401 + + return + except ImportError: + compile_helpers() + + def get_train_data_iterator() -> Iterator: if torch.distributed.get_rank() == 0: - compile_helpers() + _ensure_dataset_helpers() torch.distributed.barrier() cfg = GPTDatasetConfig( random_seed=0, @@ -104,7 +134,7 @@ def get_train_data_iterator() -> Iterator: reset_attention_mask=False, eod_mask_loss=False, tokenizer=MegatronTokenizer.from_pretrained( - metadata_path={"library": "null-text"}, + metadata_path={"library": "null"}, vocab_size=_SEQUENCE_LENGTH, ), mid_level_dataset_surplus=0.005, @@ -157,6 +187,28 @@ def _print_ce_symbols() -> None: print() +def _print_swiglu_symbols() -> None: + """Show ``bias_swiglu_impl`` bindings in the defining module and its consumers. + + The two consumer modules capture the symbol at import time, so the defining + module and each consumer are printed separately to verify all were rebound. + """ + import megatron.core.fusions.fused_bias_swiglu as defining + import megatron.core.transformer.mlp as mlp + import megatron.core.transformer.moe.shared_experts as shared + + print("\n=== Resolved SwiGLU symbols ===") + for label, mod in ( + ("fusions.fused_bias_swiglu (defines)", defining), + ("transformer.mlp (consumes)", mlp), + ("transformer.moe.shared_experts (consumes)", shared), + ): + fn = mod.bias_swiglu_impl + tag = "Liger" if getattr(fn, "__liger_patched__", False) else "Megatron" + print(f" {label:44s} → {fn.__name__} [{tag}]") + print() + + def main() -> None: # TP=1, DP=2 — CE patch (TP=1 only). Norms are correct under any TP value, so # demonstrating both Liger features in one script means running data-parallel. @@ -171,6 +223,7 @@ def main() -> None: print(gpt_model) _print_norm_classes(gpt_model) _print_ce_symbols() + _print_swiglu_symbols() ddp_cfg = DistributedDataParallelConfig( grad_reduce_in_fp32=False, diff --git a/examples/megatron/run_mode2_hand_spec.py b/examples/megatron/run_mode2_hand_spec.py index 4053a7adc..268862d64 100644 --- a/examples/megatron/run_mode2_hand_spec.py +++ b/examples/megatron/run_mode2_hand_spec.py @@ -3,8 +3,8 @@ Adapted from Megatron's ``examples/run_simple_mcore_train_loop.py``. The relevant additions (vs. that file) are: - 1. Direct imports of ``LigerMegatronRMSNorm`` and ``LigerMegatronCrossEntropy`` - (no monkey-patch). + 1. Direct imports of ``LigerMegatronRMSNorm``, ``LigerMegatronCrossEntropy`` + and ``LigerMegatronSwiGLU`` (no monkey-patch). 2. ``model_provider()`` assembles a ``TransformerBlockSubmodules`` by hand, placing ``LigerMegatronRMSNorm`` into every norm slot: @@ -20,9 +20,16 @@ a ``LigerMegatronCrossEntropy`` instance. Cross-entropy has no spec slot in Megatron, so subclassing is the symmetric "hand-built" path. - 4. ``_print_norm_classes`` + ``_print_ce_class`` after model construction - — print the resolved class for every norm slot AND the resolved CE - class on the model so you can verify Liger took over. + 4. ``_LigerSwiGLUMLP(MLP)`` owns a ``LigerMegatronSwiGLU`` instance and is + placed in the ``mlp`` spec slot. The activation likewise has no usable + spec slot of its own — ``MLPSubmodules.activation_func`` is only read when + ``config.use_te_activation_func`` is set — but the MLP *class* is a slot, + so subclassing it is the hand-built equivalent. + + 5. ``_print_norm_classes`` + ``_print_ce_class`` + ``_print_swiglu_class`` + after model construction — print the resolved class for every norm slot, + the resolved CE class, and the per-MLP SwiGLU module, so you can verify + Liger took over. Run with: torchrun --nproc_per_node=2 --master_addr=127.0.0.1 --master_port=29500 \\ @@ -38,6 +45,7 @@ class on the model so you can verify Liger took over. from typing import Iterator import torch +import torch.nn.functional as F from megatron.core import dist_checkpointing from megatron.core import parallel_state @@ -73,6 +81,7 @@ class on the model so you can verify Liger took over. # --- Liger integration: Mode 2 --------------------------------------------- from liger_kernel.megatron import LigerMegatronCrossEntropy from liger_kernel.megatron import LigerMegatronRMSNorm +from liger_kernel.megatron import LigerMegatronSwiGLU # --------------------------------------------------------------------------- @@ -81,6 +90,39 @@ class on the model so you can verify Liger took over. _NUM_LAYERS = 2 _LABEL_SMOOTHING = 0.1 +# One shared, stateless activation for every layer. +# +# Keep this as a module-level singleton, not a child module. +# If registered as a child, Megatron checkpointing calls ``sharded_state_dict()`` on it, +# but ``LigerMegatronSwiGLU`` is stateless and has nothing to shard. +_LIGER_SWIGLU = LigerMegatronSwiGLU() + + +class _LigerSwiGLUMLP(MLP): + """``MLP`` subclass whose activation is ``LigerMegatronSwiGLU``. + + SwiGLU has no reliable standalone spec slot for this path. + ``MLPSubmodules.activation_func`` is only used in the TransformerEngine path, + while standard ``MLP.forward`` resolves ``bias_swiglu_impl`` directly. + + So we use the MLP class slot: subclass ``MLP`` and call Liger SwiGLU explicitly. + ``forward`` keeps only the fused-SwiGLU branch relevant to this config. + """ + + def forward(self, hidden_states, per_token_scale=None, **kwargs): + assert per_token_scale is None, "this example does not cover the MoE per-token-scale path" + + # [s, b, 2 * ffn/p] — gate and up halves concatenated on the last dim. + intermediate, bias = self.linear_fc1(hidden_states) + intermediate = _LIGER_SWIGLU( + intermediate, + bias, + self.config.activation_func_fp8_input_store, + False, # cpu_offload_input + ) + # [s, b, h] + return self.linear_fc2(intermediate) + class _LigerCEGPTModel(GPTModel): """``GPTModel`` subclass that routes its loss through ``LigerMegatronCrossEntropy``. @@ -124,6 +166,9 @@ def model_provider() -> GPTModel: use_cpu_initialization=True, pipeline_dtype=torch.float32, normalization="RMSNorm", + gated_linear_unit=True, + activation_func=F.silu, + add_bias_linear=False, ) # ↓↓ Mode 2 — explicit slot-level placement of LigerMegatronRMSNorm ↓↓ @@ -144,8 +189,8 @@ def model_provider() -> GPTModel: ), self_attn_bda=get_bias_dropout_add, pre_mlp_layernorm=LigerMegatronRMSNorm, - mlp=partial( - MLP.as_mlp_submodule, + mlp=ModuleSpec( + module=_LigerSwiGLUMLP, submodules=MLPSubmodules( linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear, @@ -174,9 +219,23 @@ def model_provider() -> GPTModel: ) +def _ensure_dataset_helpers() -> None: + """Build Megatron's C++ dataset helpers only if not already compiled. + + ``compile_helpers()`` shells out to ``make``, which fails on pip-installed + megatron-core wheels (no Makefile, but the extension is already present). + """ + try: + from megatron.core.datasets import helpers_cpp # noqa: F401 + + return + except ImportError: + compile_helpers() + + def get_train_data_iterator() -> Iterator: if torch.distributed.get_rank() == 0: - compile_helpers() + _ensure_dataset_helpers() torch.distributed.barrier() cfg = GPTDatasetConfig( random_seed=0, @@ -185,7 +244,7 @@ def get_train_data_iterator() -> Iterator: reset_attention_mask=False, eod_mask_loss=False, tokenizer=MegatronTokenizer.from_pretrained( - metadata_path={"library": "null-text"}, + metadata_path={"library": "null"}, vocab_size=_SEQUENCE_LENGTH, ), mid_level_dataset_surplus=0.005, @@ -240,6 +299,27 @@ def _print_ce_class(model: torch.nn.Module) -> None: print() +def _print_swiglu_class(model: torch.nn.Module) -> None: + """Show that every MLP routes its activation through ``LigerMegatronSwiGLU``. + + Mode 2 does not touch Megatron's module-level ``bias_swiglu_impl`` at all, so unlike + Mode 1 there is no patched symbol to inspect. The activation is also intentionally + not a registered child (see ``_LIGER_SWIGLU``), so it will not appear in the module + tree either — the evidence is that each ``mlp`` slot resolved to ``_LigerSwiGLUMLP``. + """ + print("=== Resolved SwiGLU MLPs ===") + found = 0 + for name, mod in model.named_modules(): + if isinstance(mod, _LigerSwiGLUMLP): + print(f" {name:50s} {type(mod).__name__} → {type(_LIGER_SWIGLU).__name__}") + found += 1 + if not found: + print(" (none found — the MLP spec slot did not pick up _LigerSwiGLUMLP)") + else: + print(f" in_place={_LIGER_SWIGLU.in_place}") + print() + + def main() -> None: # TP=1, DP=2 — CE patch (TP=1 only). Norms are correct under any TP value, so # demonstrating both Liger features in one script means running data-parallel. @@ -254,6 +334,7 @@ def main() -> None: print(gpt_model) _print_norm_classes(gpt_model) _print_ce_class(gpt_model) + _print_swiglu_class(gpt_model) ddp_cfg = DistributedDataParallelConfig( grad_reduce_in_fp32=False, diff --git a/src/liger_kernel/megatron/__init__.py b/src/liger_kernel/megatron/__init__.py index 5aa77f909..d9a88f148 100644 --- a/src/liger_kernel/megatron/__init__.py +++ b/src/liger_kernel/megatron/__init__.py @@ -5,10 +5,14 @@ LayerNormBuilder protocol. LigerMegatronCrossEntropy — drop-in for Megatron-LM's vocab-parallel cross-entropy (Megatron defaults). Supports all TP sizes. + LigerMegatronSwiGLU — drop-in for Megatron-Core's ``bias_swiglu_impl``, + the fused gated-SiLU activation used by the dense MLP and the MoE + shared experts. Falls back to Megatron's implementation for bias, + FP8 input store, and CPU activation offload. apply_liger_kernel_to_megatron — patches Megatron-Core so existing training scripts pick up Liger kernels with one line. Currently supports - RMSNorm (via BackendSpecProvider) plus both the fused and unfused - vocab-parallel cross-entropy paths. + RMSNorm (via BackendSpecProvider), both the fused and unfused + vocab-parallel cross-entropy paths, and SwiGLU. The general-purpose ``LigerVocabParallelCrossEntropy`` Module lives under ``liger_kernel.transformers`` alongside the other nn.Module wrappers; the @@ -20,9 +24,11 @@ from liger_kernel.megatron.cross_entropy import LigerMegatronCrossEntropy from liger_kernel.megatron.monkey_patch import apply_liger_kernel_to_megatron from liger_kernel.megatron.rms_norm import LigerMegatronRMSNorm +from liger_kernel.megatron.swiglu import LigerMegatronSwiGLU __all__ = [ "LigerMegatronCrossEntropy", "LigerMegatronRMSNorm", + "LigerMegatronSwiGLU", "apply_liger_kernel_to_megatron", ] diff --git a/src/liger_kernel/megatron/monkey_patch.py b/src/liger_kernel/megatron/monkey_patch.py index aebce69af..fd8de5f7d 100644 --- a/src/liger_kernel/megatron/monkey_patch.py +++ b/src/liger_kernel/megatron/monkey_patch.py @@ -2,6 +2,7 @@ from __future__ import annotations +import inspect import logging logger = logging.getLogger(__name__) @@ -12,6 +13,7 @@ def apply_liger_kernel_to_megatron( rms_norm: bool = True, cross_entropy: bool = False, + swiglu: bool = False, ) -> None: """Patch Megatron-Core to use Liger Triton kernels. @@ -38,6 +40,15 @@ def apply_liger_kernel_to_megatron( wrapper additionally honors a runtime ``label_smoothing`` argument, matching native's ``(logits, target, label_smoothing=0.0, tp_group=None)``. + swiglu: When ``True`` replace + ``megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl`` with Liger's + Triton SiLU-multiply kernel. Covers the dense ``MLP`` and the MoE + ``SharedExpertMLP``, i.e. every site Megatron reaches when + ``config.bias_activation_fusion=True``, ``config.gated_linear_unit=True`` + and ``config.activation_func is F.silu``. Default ``False`` so adopters + opt in explicitly. Calls Liger cannot serve (a non-``None`` bias, + ``activation_func_fp8_input_store=True``, or CPU activation offloading) + transparently defer to Megatron's original implementation. Notes: Call this BEFORE building your model. Patching after instantiation @@ -61,6 +72,8 @@ def apply_liger_kernel_to_megatron( if cross_entropy: _patch_fused_vocab_parallel_cross_entropy() _patch_vocab_parallel_cross_entropy() + if swiglu: + _patch_bias_swiglu_impl() def _patch_local_spec_provider_layer_norm() -> None: @@ -72,16 +85,20 @@ def _patch_local_spec_provider_layer_norm() -> None: return # already patched original_layer_norm = backends.LocalSpecProvider.layer_norm - - def patched_layer_norm( - self, - rms_norm: bool = False, - for_qk: bool = False, - has_residual: bool = False, - ): - if rms_norm: + _original_sig = inspect.signature(original_layer_norm) + + def patched_layer_norm(self, *args, **kwargs): + # Forward through the *original* signature rather than restating it. Megatron has + # changed this method's parameters across releases (``has_residual`` was added + # after ``for_qk``), and a hardcoded forwarding call raises TypeError on any + # version whose signature does not match exactly. Binding to the real signature + # also means ``rms_norm`` is read correctly whether the caller passed it + # positionally or by keyword. + bound = _original_sig.bind(self, *args, **kwargs) + bound.apply_defaults() + if bound.arguments.get("rms_norm", False): return LigerMegatronRMSNorm - return original_layer_norm(self, rms_norm=rms_norm, for_qk=for_qk, has_residual=has_residual) + return original_layer_norm(*bound.args, **bound.kwargs) setattr(patched_layer_norm, _PATCH_MARKER, True) setattr(patched_layer_norm, "__wrapped__", original_layer_norm) @@ -260,3 +277,107 @@ def liger_vocab_parallel_cross_entropy( logger.info( "Patched megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy with Liger cross-entropy." ) + + +# Modules that do ``from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl`` +# and therefore hold their own binding. Must be rebound individually after patching. +# Verified against Megatron-LM ``main``; missing entries are skipped. +_BIAS_SWIGLU_CONSUMER_MODULES = ( + "megatron.core.transformer.mlp", + "megatron.core.transformer.moe.shared_experts", +) + + +def _patch_bias_swiglu_impl() -> None: + """Replace ``megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl`` with Liger. + + Covers ``MLP.forward`` and ``SharedExpertMLP.forward`` when + ``config.bias_activation_fusion=True``, ``config.gated_linear_unit=True`` + and ``config.activation_func is F.silu``. + + Not patched: ``weighted_bias_swiglu_impl`` (needs routing-weight grad), + the unfused ``glu()`` closure (no module-level symbol), and + ``config.use_te_activation_func=True`` (TE owns that path). + + Unsupported configs (bias, FP8, CPU offload) defer to Megatron's original at call time. + """ + try: + import megatron.core.fusions.fused_bias_swiglu as fused_swiglu + except ImportError as exc: + raise ImportError( + "apply_liger_kernel_to_megatron(swiglu=True) requires megatron-core to be " + "installed. Expected symbol path: " + "megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl." + ) from exc + + if not hasattr(fused_swiglu, "bias_swiglu_impl"): + raise ImportError( + "megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl not found. The symbol " + "path may have changed in your Megatron-LM version. Please file an issue on " + "https://github.com/linkedin/Liger-Kernel with your megatron-core version." + ) + + if getattr(fused_swiglu.bias_swiglu_impl, _PATCH_MARKER, False): + # Already patched -- re-run consumer rebinding for modules imported after the + # first patch call (they'd still hold Megatron's original binding). + liger_impl = fused_swiglu.bias_swiglu_impl + _rebind_consumer_symbol( + _BIAS_SWIGLU_CONSUMER_MODULES, + "bias_swiglu_impl", + liger_impl.__wrapped__, + liger_impl, + ) + return + + original = fused_swiglu.bias_swiglu_impl + + from liger_kernel.megatron.swiglu import LigerMegatronSwiGLU + + # Reuse a single instance: stateless except for the fallback-log dedup set. + swiglu_module = LigerMegatronSwiGLU(fallback_impl=original) + + def liger_bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False): + return swiglu_module(input, bias, fp8_input_store, cpu_offload_input) + + setattr(liger_bias_swiglu_impl, _PATCH_MARKER, True) + setattr(liger_bias_swiglu_impl, "__wrapped__", original) + fused_swiglu.bias_swiglu_impl = liger_bias_swiglu_impl + + patched_consumers = _rebind_consumer_symbol( + _BIAS_SWIGLU_CONSUMER_MODULES, + "bias_swiglu_impl", + original, + liger_bias_swiglu_impl, + ) + + logger.info( + "Patched megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl with Liger SwiGLU (also rebound in: %s).", + ", ".join(patched_consumers) if patched_consumers else "no already-imported consumers", + ) + + +def _rebind_consumer_symbol( + module_names, + attr_name: str, + original, + replacement, +): + """Point already-imported ``from X import attr`` bindings at ``replacement``. + + Only rebinds modules already in ``sys.modules`` whose binding is still ``original``. + Modules imported later pick up the patched symbol naturally. + + Returns the list of module names actually rebound. + """ + import sys + + rebound = [] + for name in module_names: + module = sys.modules.get(name) + if module is None: + continue # not imported yet + if getattr(module, attr_name, None) is not original: + continue # absent, or already pointing somewhere else + setattr(module, attr_name, replacement) + rebound.append(name) + return rebound diff --git a/src/liger_kernel/megatron/swiglu.py b/src/liger_kernel/megatron/swiglu.py new file mode 100644 index 000000000..7651ef8af --- /dev/null +++ b/src/liger_kernel/megatron/swiglu.py @@ -0,0 +1,162 @@ +"""Megatron-Core compatible SwiGLU backed by the Liger Triton SiLU-multiply kernel. + +Drop-in for ``megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl``, which is the +symbol ``MLP.forward`` and ``SharedExpertMLP.forward`` call when +``config.bias_activation_fusion=True``, ``config.gated_linear_unit=True`` and +``config.activation_func is F.silu``. + +Megatron's native implementation is a chain of ``@jit_fuser`` (TorchScript) helpers:: + + y_1, y_2 = torch.chunk(y, 2, -1) + return F.silu(y_1) * y_2 + +Liger fuses the whole thing into a single Triton kernel that computes sigmoid / silu / +multiply in registers, and recomputes ``silu`` in the backward instead of saving it. + +This uses ``LigerFusedGateUpSiLUMulFunction`` rather than the two-tensor +``LigerSiLUMulFunction``, because Megatron hands over gate and up concatenated in one +buffer. Splitting them with ``torch.chunk`` yields non-contiguous views that +``@ensure_contiguous`` materializes into two full-size copies, plus a ``torch.cat`` to +reassemble the gradient. Measured on H100 at ``[2048*4, 2*32768]`` bf16, that bridge ran +3.4x slower than the fused kernel and 2.1x slower than Megatron's own TorchScript path -- +a regression dressed up as an optimization. The fused kernel reads both halves from the +single buffer via a column offset instead. + +Unsupported configurations (bias, FP8 input store, CPU activation offload) transparently +fall back to the native implementation, so enabling Liger can never change behavior for +those users. +""" + +from __future__ import annotations + +import logging + +from typing import Callable +from typing import Optional + +import torch +import torch.nn as nn + +# Force-import the submodule so liger_kernel.ops.swiglu can resolve +# torch.distributed.tensor.DTensor on torch 2.11+, where the subpackage is no longer +# auto-loaded as an attribute of torch.distributed. +import torch.distributed.tensor # noqa: F401 # isort: skip + +from liger_kernel.ops import LigerFusedGateUpSiLUMulFunction + +logger = logging.getLogger(__name__) + + +def _unsupported_reason( + bias: Optional[torch.Tensor], + fp8_input_store: bool, + cpu_offload_input: bool, +) -> Optional[str]: + """Return a non-empty reason string if Liger cannot serve this call, else None.""" + if bias is not None: + # Liger's kernel has no bias term; fusing it is tracked as follow-up work. + return "add_bias_linear=True (bias is not None)" + if fp8_input_store: + # Megatron stores activations as fp8 and restores dtype in backward; + # incompatible with Liger's in-place gradient writes. + return "config.activation_func_fp8_input_store=True" + if cpu_offload_input: + # In-place gradient writes into a CPU-offloaded buffer are unvalidated. + return "CPU activation offloading enabled" + return None + + +class LigerMegatronSwiGLU(nn.Module): + """``bias_swiglu_impl``-compatible SwiGLU using Liger's Triton kernel. + + The call signature matches Megatron's exactly -- including positional order, since + both ``MLP.forward`` and ``SharedExpertMLP.forward`` call it positionally:: + + bias_swiglu_impl(input, bias, fp8_input_store, cpu_offload_input) + + Args: + fallback_impl: Callable used for configurations Liger cannot serve (bias, FP8 + input store, CPU offload). The monkey patch passes Megatron's captured + original here. When ``None`` (direct Mode 2 construction) an unsupported + configuration raises instead of silently changing numerics. + in_place: Write the backward gradient into the fc1 output buffer instead of + allocating a new one. Saves one activation-sized allocation -- measured on + H100 at ``ffn_local=32768`` it moves peak memory from 20% above Megatron to + 20% below it -- at the cost of destroying the fc1 output during backward. + + Default False, matching ``LigerMegatronRMSNorm``. In the standard dense-MLP + path nothing reads the fc1 output after the activation (``linear_fc1``'s + backward needs its *input*, and ``linear_fc2``'s needs the activation + *output*), so this is expected to be safe; it is opt-in because that has not + been validated against every recompute / offload / CUDA-graph combination. + + Incompatible with backpropagating through the same graph twice + (``retain_graph=True`` followed by a second ``backward``, or double-backward). + That is detected and raises; see ``LigerFusedGateUpSiLUMulFunction.backward`` + for why autograd's own version-counter check cannot catch it. + """ + + def __init__(self, fallback_impl: Optional[Callable] = None, in_place: bool = False): + super().__init__() + self.fallback_impl = fallback_impl + self.in_place = in_place + # Reasons already logged, so a fallback taken every microbatch logs once, not + # once per step. + self._logged_fallbacks: set[str] = set() + + def _fallback( + self, + reason: str, + input: torch.Tensor, + bias: Optional[torch.Tensor], + fp8_input_store: bool, + cpu_offload_input: bool, + ) -> torch.Tensor: + if self.fallback_impl is None: + raise RuntimeError( + f"LigerMegatronSwiGLU cannot serve this call: {reason}. Liger's SwiGLU " + "kernel supports neither a fused bias term, FP8 input storage, nor CPU " + "activation offloading. Either construct this module with " + "`fallback_impl=megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl` " + "so unsupported calls defer to Megatron, or use " + "`apply_liger_kernel_to_megatron(swiglu=True)`, which wires that fallback " + "up automatically." + ) + if reason not in self._logged_fallbacks: + self._logged_fallbacks.add(reason) + logger.info( + "Liger SwiGLU is falling back to Megatron's native bias_swiglu_impl: %s. " + "Numerics and memory behavior are unchanged for this configuration.", + reason, + ) + return self.fallback_impl(input, bias, fp8_input_store, cpu_offload_input) + + def forward( + self, + input: torch.Tensor, + bias: Optional[torch.Tensor] = None, + fp8_input_store: bool = False, + cpu_offload_input: bool = False, + ) -> torch.Tensor: + reason = _unsupported_reason(bias, fp8_input_store, cpu_offload_input) + if reason is not None: + return self._fallback(reason, input, bias, fp8_input_store, cpu_offload_input) + + ori_shape = input.shape + if len(ori_shape) not in (2, 3): + raise ValueError( + f"LigerMegatronSwiGLU expects a 2D [tokens, 2*ffn] or 3D [seq, batch, 2*ffn] " + f"input, matching Megatron's bias_swiglu_impl; got shape {tuple(ori_shape)}." + ) + if ori_shape[-1] % 2 != 0: + raise ValueError( + "LigerMegatronSwiGLU expects the gate and up projections concatenated along " + f"the last dimension, so it must be even; got {ori_shape[-1]}. This usually " + "means config.gated_linear_unit is False, in which case Megatron does not " + "route through bias_swiglu_impl at all." + ) + + # Mirror Megatron's own reshape-compute-restore so 3D activations behave + # identically. + output = LigerFusedGateUpSiLUMulFunction.apply(input.view(-1, ori_shape[-1]), self.in_place) + return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) diff --git a/src/liger_kernel/ops/__init__.py b/src/liger_kernel/ops/__init__.py index f6676c07e..32876b6e4 100644 --- a/src/liger_kernel/ops/__init__.py +++ b/src/liger_kernel/ops/__init__.py @@ -88,6 +88,7 @@ from liger_kernel.ops.rope import rope_forward # noqa: F401 from liger_kernel.ops.softmax import LigerSoftmaxFunction # noqa: F401 from liger_kernel.ops.sparsemax import LigerSparsemaxFunction # noqa: F401 +from liger_kernel.ops.swiglu import LigerFusedGateUpSiLUMulFunction # noqa: F401 from liger_kernel.ops.swiglu import LigerSiLUMulFunction # noqa: F401 from liger_kernel.ops.swiglu import swiglu_backward # noqa: F401 from liger_kernel.ops.swiglu import swiglu_forward # noqa: F401 diff --git a/src/liger_kernel/ops/swiglu.py b/src/liger_kernel/ops/swiglu.py index 047832a5c..c1cf50d6f 100644 --- a/src/liger_kernel/ops/swiglu.py +++ b/src/liger_kernel/ops/swiglu.py @@ -281,3 +281,148 @@ def backward(ctx, dc): dc = dc * down_multiplier a, b = swiglu_backward(a, b, dc, gate_multiplier) return a, b, None, None + + +# --------------------------------------------------------------------------- +# Fused gate-up variant +# --------------------------------------------------------------------------- +# For fused ``[tokens, 2*n]`` gate-up tensors (Megatron, HF ``gate_up_proj``): the +# kernels read both halves via a column offset into the single buffer -- no copies, +# no cat. Input row stride is ``2*n``, output is ``n``. + + +@triton.jit +def _swiglu_fused_gate_up_forward_kernel( + y_ptr, c_ptr, in_stride, out_stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr +): + program_id = tl.program_id(0).to(tl.int64) + + y_ptr += program_id * in_stride + c_ptr += program_id * out_stride + + col_offsets = tl.arange(0, BLOCK_SIZE) + mask = col_offsets < n_cols + + # Gate occupies columns [0, n), up occupies [n, 2n) of the same row. + gate = tl.load(y_ptr + col_offsets, mask=mask, other=0).to(tl.float32) + up = tl.load(y_ptr + n_cols + col_offsets, mask=mask, other=0) + tl.store(c_ptr + col_offsets, silu(gate).cast(up.dtype) * up, mask=mask) + + +@triton.jit +def _swiglu_fused_gate_up_backward_kernel( + dc_ptr, y_ptr, dy_ptr, in_stride, out_stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr +): + program_id = tl.program_id(0).to(tl.int64) + + dc_ptr += program_id * out_stride + y_ptr += program_id * in_stride + dy_ptr += program_id * in_stride + + col_offsets = tl.arange(0, BLOCK_SIZE) + mask = col_offsets < n_cols + + dc = tl.load(dc_ptr + col_offsets, mask=mask, other=0) + gate = tl.load(y_ptr + col_offsets, mask=mask, other=0).to(tl.float32) + up = tl.load(y_ptr + n_cols + col_offsets, mask=mask, other=0) + + # Recompute silu from saved input. When dy_ptr aliases y_ptr (in_place=True), all + # loads precede all stores and each program owns one row, so aliasing is safe here. + # Caller-level safety is enforced in swiglu_fused_gate_up_backward. + sig = tl.sigmoid(gate) + silu_gate = gate * sig + d_gate = dc * (silu_gate * (1 - sig) + sig) * up + d_up = dc * silu_gate + + tl.store(dy_ptr + col_offsets, d_gate, mask=mask) + tl.store(dy_ptr + n_cols + col_offsets, d_up, mask=mask) + + +def swiglu_fused_gate_up_forward(y): + """SwiGLU over a fused ``[..., 2 * n]`` gate-up tensor. Returns ``(y, c)``.""" + ori_shape = y.shape + two_n = ori_shape[-1] + if two_n % 2 != 0: + raise ValueError(f"fused gate-up input must have an even trailing dim; got {two_n}.") + n_cols = two_n // 2 + + y = y.view(-1, two_n) + n_rows = y.shape[0] + c = torch.empty(n_rows, n_cols, dtype=y.dtype, device=y.device) + + BLOCK_SIZE, num_warps = calculate_settings(n_cols) + _swiglu_fused_gate_up_forward_kernel[(n_rows,)]( + y, + c, + y.stride(-2), + c.stride(-2), + n_cols=n_cols, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + ) + return y, c.view(*ori_shape[:-1], n_cols) + + +def swiglu_fused_gate_up_backward(y, dc, in_place=False): + """Gradient w.r.t. the fused ``[..., 2 * n]`` gate-up tensor. + + Args: + in_place: Overwrite ``y`` with the gradient in place instead of allocating a new + ``[..., 2*n]`` buffer. Saves ~1 GB peak on H100 at ``[2048*4, 2*32768]`` bf16. + Default False -- Megatron may hold references to the fc1 output for activation + recompute or CUDA-graph capture, so clobbering it requires care. + """ + two_n = y.shape[-1] + n_cols = two_n // 2 + + y = y.view(-1, two_n) + dc = dc.view(-1, n_cols) + n_rows = dc.shape[0] + dy = y if in_place else torch.empty_like(y) + + BLOCK_SIZE, num_warps = calculate_settings(n_cols) + _swiglu_fused_gate_up_backward_kernel[(n_rows,)]( + dc, + y, + dy, + y.stride(-2), + dc.stride(-2), + n_cols=n_cols, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + ) + return dy + + +class LigerFusedGateUpSiLUMulFunction(torch.autograd.Function): + """SwiGLU for a single tensor holding gate and up concatenated on the last dim. + + Equivalent to ``LigerSiLUMulFunction.apply(*torch.chunk(y, 2, -1))`` but without the + contiguous copies that chunking forces, and without concatenating the gradient. Used by + the Megatron integration, where ``linear_fc1`` always emits this layout. + """ + + @staticmethod + @ensure_contiguous + def forward(ctx, y, in_place=False): + y, c = swiglu_fused_gate_up_forward(y) + ctx.save_for_backward(y) + ctx.in_place = in_place + ctx.already_backward = False + return c + + @staticmethod + @ensure_contiguous + def backward(ctx, dc): + (y,) = ctx.saved_tensors + if ctx.in_place: + # Triton writes bypass autograd's version counter, so a second backward would + # silently read gradients instead of activations and produce wrong results. + if ctx.already_backward: + raise RuntimeError( + "LigerFusedGateUpSiLUMulFunction(in_place=True) can only run backward " + "once -- the saved activation is overwritten with the gradient. " + "Use in_place=False for retain_graph=True or double-backward." + ) + ctx.already_backward = True + return swiglu_fused_gate_up_backward(y, dc, ctx.in_place), None diff --git a/test/megatron/test_monkey_patch.py b/test/megatron/test_monkey_patch.py index 391c1176d..6bef010e2 100644 --- a/test/megatron/test_monkey_patch.py +++ b/test/megatron/test_monkey_patch.py @@ -20,6 +20,7 @@ 3. RMSNorm patch tests 4. Cross-kernel public-API surface tests 5. End-to-end integration through patched CE symbols + 6. SwiGLU patch tests """ import sys @@ -59,6 +60,7 @@ def _install_fake_megatron_ce( Returns a tuple ``(fused_ce_module, unfused_ce_module)`` so tests can inspect what the patch helpers wrote onto them. """ + _snapshot_real_megatron() _, megatron_core = _ensure_megatron_roots() fusions = types.ModuleType("megatron.core.fusions") fused_ce = types.ModuleType("megatron.core.fusions.fused_cross_entropy") @@ -119,6 +121,7 @@ def _install_fake_megatron_rms_norm( level patch only displaces that fallback; set False to verify the no-op path taken under TE / Apex. """ + _snapshot_real_megatron() _, megatron_core = _ensure_megatron_roots() backends = None @@ -188,26 +191,73 @@ class _SomeOtherNorm: return backends, transformer_block +_SAVED_REAL_MEGATRON: dict = {} +_STUBS_ACTIVE = False + + +def _is_megatron_module(name: str) -> bool: + return name == "megatron" or name.startswith("megatron.") + + +def _snapshot_real_megatron(): + """Record the entire ``megatron*`` region of ``sys.modules`` before stubbing it. + + Only has an effect when megatron-core is genuinely installed, e.g. an environment set + up to run the end-to-end section of ``test_swiglu.py``. See ``_uninstall_fake_megatron`` + for why this has to cover the whole package rather than the specific names the stubs + replace. + """ + global _STUBS_ACTIVE + if _STUBS_ACTIVE: + return # a stub install is already in effect; don't snapshot the stubs themselves + _SAVED_REAL_MEGATRON.clear() + _SAVED_REAL_MEGATRON.update({name: mod for name, mod in sys.modules.items() if _is_megatron_module(name)}) + _STUBS_ACTIVE = True + + def _uninstall_fake_megatron(): - """Tear down every stub module installed by either installer.""" - for mod in [ - # CE side - "megatron.core.parallel_state", - "megatron.core.fusions.fused_cross_entropy", - "megatron.core.fusions", - "megatron.core.tensor_parallel.cross_entropy", - "megatron.core.tensor_parallel", - # RMSNorm side - "megatron.core.models.backends", - "megatron.core.models", - "megatron.core.transformer.transformer_block", - "megatron.core.transformer.torch_norm", - "megatron.core.transformer", - # Shared roots - "megatron.core", - "megatron", - ]: - sys.modules.pop(mod, None) + """Restore ``sys.modules`` to exactly its pre-stub state. + + Two subtleties, both of which only bite when a real megatron-core is installed: + + 1. Popping a real module and walking away is not neutral. The next import rebuilds it + as a *new* object with fresh globals, silently discarding Megatron's process-global + parallel state -- so a later ``initialize_model_parallel`` appears not to have run. + 2. Evicting a hand-picked list of names is not enough. Submodules outside the list + (e.g. ``megatron.core.tensor_parallel.random``) survive holding references to the + evicted parent, so the process ends up with two live ``parallel_state`` objects and + whichever one gets initialized is not the one that gets read. + + Clearing and restoring the whole ``megatron*`` prefix avoids both. + """ + global _STUBS_ACTIVE + if not _STUBS_ACTIVE: + return # nothing was stubbed; leave a real install alone + + stubbed = [n for n in list(sys.modules) if _is_megatron_module(n)] + for name in stubbed: + del sys.modules[name] + sys.modules.update(_SAVED_REAL_MEGATRON) + + # sys.modules is not the only place a stub lands. Importing `megatron.core.x` also + # binds `x` as an *attribute* of the `megatron.core` package object, and restoring + # sys.modules does not undo that binding. A later `from megatron.core import x` reads + # the attribute and gets the stub straight back -- leaving two live copies of a module + # whose globals are process-wide state, which is how `initialize_model_parallel` ends + # up writing to one copy while everything else reads the other. + for name in stubbed: + parent_name, _, child = name.rpartition(".") + parent = sys.modules.get(parent_name) + if parent is None: + continue + real = _SAVED_REAL_MEGATRON.get(name) + if real is not None: + setattr(parent, child, real) + elif hasattr(parent, child): + delattr(parent, child) + + _SAVED_REAL_MEGATRON.clear() + _STUBS_ACTIVE = False @pytest.fixture @@ -228,6 +278,96 @@ def fake_megatron_rms_norm(): _uninstall_fake_megatron() +def _install_fake_megatron_swiglu( + with_symbol: bool = True, + with_consumer_modules: bool = True, +): + """Install the SwiGLU slice of the Megatron stub. + + Returns ``(fused_swiglu_module, consumer_modules_dict)``. + + The consumer modules matter here in a way they don't for CE or RMSNorm: real Megatron + does ``from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl`` in + ``transformer/mlp.py`` and ``transformer/moe/shared_experts.py``, so each holds its own + binding. A patch that only rewrites the defining module leaves every call site on + native Megatron. The stub reproduces those by-name imports so the tests can catch that. + """ + _snapshot_real_megatron() + _, megatron_core = _ensure_megatron_roots() + + # ``megatron.core.fusions`` may already exist if the CE installer ran first. + fusions = sys.modules.get("megatron.core.fusions") or types.ModuleType("megatron.core.fusions") + fused_swiglu = types.ModuleType("megatron.core.fusions.fused_bias_swiglu") + + def original_bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False): + """Stub of Megatron's native impl. Returns Megatron's reference result so + fallback tests can assert numerics, and records that it ran.""" + import torch + import torch.nn.functional as F + + original_bias_swiglu_impl.call_count += 1 + y = input if bias is None else input + bias + y_1, y_2 = torch.chunk(y, 2, -1) + return F.silu(y_1) * y_2 + + original_bias_swiglu_impl.call_count = 0 + + if with_symbol: + fused_swiglu.bias_swiglu_impl = original_bias_swiglu_impl + # The MoE-weighted variant is deliberately out of scope for the current patch; it is + # present on the stub so a test can assert Liger leaves it alone. + fused_swiglu.weighted_bias_swiglu_impl = lambda *a, **kw: None + + sys.modules["megatron.core.fusions"] = fusions + sys.modules["megatron.core.fusions.fused_bias_swiglu"] = fused_swiglu + megatron_core.fusions = fusions + fusions.fused_bias_swiglu = fused_swiglu + + consumers = {} + if with_consumer_modules: + transformer = sys.modules.get("megatron.core.transformer") or types.ModuleType("megatron.core.transformer") + moe = types.ModuleType("megatron.core.transformer.moe") + mlp = types.ModuleType("megatron.core.transformer.mlp") + shared_experts = types.ModuleType("megatron.core.transformer.moe.shared_experts") + + # Reproduce the by-name import: each consumer gets its own binding to the + # original function object. + if with_symbol: + mlp.bias_swiglu_impl = original_bias_swiglu_impl + shared_experts.bias_swiglu_impl = original_bias_swiglu_impl + + sys.modules["megatron.core.transformer"] = transformer + sys.modules["megatron.core.transformer.moe"] = moe + sys.modules["megatron.core.transformer.mlp"] = mlp + sys.modules["megatron.core.transformer.moe.shared_experts"] = shared_experts + megatron_core.transformer = transformer + transformer.mlp = mlp + transformer.moe = moe + moe.shared_experts = shared_experts + + consumers = {"mlp": mlp, "shared_experts": shared_experts} + + return fused_swiglu, consumers + + +@pytest.fixture +def fake_megatron_swiglu(): + fused_swiglu, consumers = _install_fake_megatron_swiglu() + try: + yield fused_swiglu, consumers + finally: + _uninstall_fake_megatron() + + +@pytest.fixture +def fake_megatron_swiglu_no_consumers(): + fused_swiglu, _ = _install_fake_megatron_swiglu(with_consumer_modules=False) + try: + yield fused_swiglu + finally: + _uninstall_fake_megatron() + + # =========================================================================== # 2. Cross-entropy patch tests # =========================================================================== @@ -793,6 +933,7 @@ def test_import_from_root(): try: from liger_kernel.megatron import LigerMegatronCrossEntropy # noqa: F401 from liger_kernel.megatron import LigerMegatronRMSNorm # noqa: F401 + from liger_kernel.megatron import LigerMegatronSwiGLU # noqa: F401 from liger_kernel.megatron import apply_liger_kernel_to_megatron # noqa: F401 except Exception: pytest.fail("Importing public Megatron symbols from liger_kernel.megatron failed.") @@ -1023,3 +1164,288 @@ def test_rms_norm_only_patch_does_not_touch_ce_symbols(fake_megatron_ce): assert fused_ce.fused_vocab_parallel_cross_entropy is fused_before assert unfused_ce.vocab_parallel_cross_entropy is unfused_before + + +# =========================================================================== +# 6. SwiGLU patch tests +# =========================================================================== +# Liger replaces ``megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl`` — the symbol +# ``MLP.forward`` and ``SharedExpertMLP.forward`` call when bias_activation_fusion is on. +# Unlike CE and RMSNorm, its consumers use by-name imports, so the patch has to rewrite +# their module attributes too; several tests below exist specifically to enforce that. + +from liger_kernel.utils import infer_device # noqa: E402 + +# Triton needs a real accelerator. Patch-mechanism tests all run on CPU; only the +# numerical end-to-end ones are gated. +_requires_accelerator = pytest.mark.skipif( + infer_device() == "cpu", + reason="Triton kernels require an accelerator", +) + + +def _real_megatron_core_installed() -> bool: + """True when a genuine megatron-core is importable, as opposed to this file's stub. + + Evaluated once at import time, before any test installs a stub into ``sys.modules``. + The "missing megatron-core" tests can only assert anything meaningful when the package + really is absent -- which is the assumption in this repo's CI, but not in an + environment where someone installed megatron-core to run the e2e suite. + """ + import importlib.util + + try: + return importlib.util.find_spec("megatron.core.fusions.fused_bias_swiglu") is not None + except (ImportError, ValueError): + return False + + +_REAL_MEGATRON = _real_megatron_core_installed() + + +# --------------------------------------------------------------------------- +# 6.1 Symbol replacement + idempotency + opt-in. +# --------------------------------------------------------------------------- + + +def test_swiglu_patch_replaces_symbol(fake_megatron_swiglu): + fused_swiglu, _ = fake_megatron_swiglu + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + original = fused_swiglu.bias_swiglu_impl + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + + assert fused_swiglu.bias_swiglu_impl is not original + assert fused_swiglu.bias_swiglu_impl.__name__ == "liger_bias_swiglu_impl" + assert fused_swiglu.bias_swiglu_impl.__wrapped__ is original + + +def test_swiglu_patch_is_opt_in(fake_megatron_swiglu): + """``swiglu`` defaults to False — the symbol must be untouched unless asked for.""" + fused_swiglu, consumers = fake_megatron_swiglu + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + before = fused_swiglu.bias_swiglu_impl + apply_liger_kernel_to_megatron(rms_norm=False) + + assert fused_swiglu.bias_swiglu_impl is before + assert consumers["mlp"].bias_swiglu_impl is before + + +def test_swiglu_patch_is_idempotent(fake_megatron_swiglu): + """Calling apply twice must not stack wrappers — the sentinel attribute guards it.""" + fused_swiglu, consumers = fake_megatron_swiglu + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + first = fused_swiglu.bias_swiglu_impl + first_mlp = consumers["mlp"].bias_swiglu_impl + + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + + assert fused_swiglu.bias_swiglu_impl is first + assert consumers["mlp"].bias_swiglu_impl is first_mlp + # And the fallback chain must still be one deep, not wrapped around itself. + assert not hasattr(fused_swiglu.bias_swiglu_impl.__wrapped__, "__wrapped__") + + +def test_swiglu_patch_leaves_weighted_variant_alone(fake_megatron_swiglu): + """The MoE token-weighted variant is explicitly out of scope for this patch; it needs + a routing-weights gradient Liger has no kernel for.""" + fused_swiglu, _ = fake_megatron_swiglu + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + before = fused_swiglu.weighted_bias_swiglu_impl + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + + assert fused_swiglu.weighted_bias_swiglu_impl is before + + +# --------------------------------------------------------------------------- +# 6.2 Consumer-module rebinding — the failure mode unique to this patch. +# --------------------------------------------------------------------------- + + +def test_swiglu_patch_rebinds_already_imported_consumers(fake_megatron_swiglu): + """``mlp.py`` and ``shared_experts.py`` do ``from ... import bias_swiglu_impl``. + + If the patch only rewrites the defining module, both keep calling native Megatron and + the patch silently does nothing. This is the single most important SwiGLU-specific + assertion in this file. + """ + fused_swiglu, consumers = fake_megatron_swiglu + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + + patched = fused_swiglu.bias_swiglu_impl + assert consumers["mlp"].bias_swiglu_impl is patched + assert consumers["shared_experts"].bias_swiglu_impl is patched + + +def test_swiglu_patch_tolerates_absent_consumer_modules(fake_megatron_swiglu_no_consumers): + """Consumers not yet imported are fine — they'll import the already-patched symbol. + The patch must not crash trying to rebind them.""" + fused_swiglu = fake_megatron_swiglu_no_consumers + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + + assert fused_swiglu.bias_swiglu_impl.__name__ == "liger_bias_swiglu_impl" + + +def test_swiglu_patch_does_not_clobber_unrelated_consumer_binding(fake_megatron_swiglu): + """If a consumer's binding isn't the original (a user already replaced it), leave it.""" + fused_swiglu, consumers = fake_megatron_swiglu + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + sentinel = object() + consumers["mlp"].bias_swiglu_impl = sentinel + + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + + assert consumers["mlp"].bias_swiglu_impl is sentinel + # The other consumer, still on the original, is rebound normally. + assert consumers["shared_experts"].bias_swiglu_impl is fused_swiglu.bias_swiglu_impl + + +# --------------------------------------------------------------------------- +# 6.3 Missing megatron-core / missing symbol → actionable ImportError. +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(_REAL_MEGATRON, reason="a real megatron-core is installed, so the import cannot fail") +def test_swiglu_patch_raises_when_megatron_missing(): + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + _uninstall_fake_megatron() + with pytest.raises(ImportError, match="fused_bias_swiglu"): + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + + +def test_swiglu_patch_raises_when_symbol_missing(): + """A renamed symbol in a future megatron-core must fail loudly, not silently no-op.""" + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + _install_fake_megatron_swiglu(with_symbol=False) + try: + with pytest.raises(ImportError, match="bias_swiglu_impl not found"): + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + finally: + _uninstall_fake_megatron() + + +# --------------------------------------------------------------------------- +# 6.4 End-to-end through the patched symbol. +# --------------------------------------------------------------------------- + + +def _swiglu_reference(y): + import torch + import torch.nn.functional as F + + y_1, y_2 = torch.chunk(y, 2, -1) + return F.silu(y_1) * y_2 + + +@_requires_accelerator +def test_patched_swiglu_symbol_computes_correct_output(fake_megatron_swiglu): + import torch + + from liger_kernel.megatron import apply_liger_kernel_to_megatron + from liger_kernel.utils import infer_device + + fused_swiglu, consumers = fake_megatron_swiglu + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + + x = torch.randn(8, 128, device=infer_device(), dtype=torch.float32) + # Call through the *consumer's* binding, exactly as MLP.forward does — positionally. + out = consumers["mlp"].bias_swiglu_impl(x, None, False, False) + + torch.testing.assert_close(out, _swiglu_reference(x), atol=1e-5, rtol=1e-5) + + +@_requires_accelerator +def test_patched_swiglu_symbol_preserves_gradients(fake_megatron_swiglu): + import torch + + from liger_kernel.megatron import apply_liger_kernel_to_megatron + from liger_kernel.utils import infer_device + + _, consumers = fake_megatron_swiglu + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + + base = torch.randn(8, 128, device=infer_device(), dtype=torch.float32) + h_ref = base.clone().requires_grad_(True) + h_liger = base.clone().requires_grad_(True) + do = torch.randn(8, 64, device=infer_device(), dtype=torch.float32) + + _swiglu_reference(h_ref).backward(do) + consumers["mlp"].bias_swiglu_impl(h_liger, None, False, False).backward(do) + + torch.testing.assert_close(h_liger.grad, h_ref.grad, atol=1e-5, rtol=1e-5) + + +def test_patched_swiglu_symbol_falls_back_for_bias(fake_megatron_swiglu): + """Bias is added pre-activation in Megatron and Liger has no bias term, so the patch + must hand the call straight back to the original. Runs on CPU — no Triton involved, + which is the point: the fallback path never touches the kernel.""" + import torch + + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + fused_swiglu, consumers = fake_megatron_swiglu + original = fused_swiglu.bias_swiglu_impl + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + + x = torch.randn(8, 128) + bias = torch.randn(128) + out = consumers["mlp"].bias_swiglu_impl(x, bias, False, False) + + assert original.call_count == 1, "expected the native impl to be invoked for the bias path" + torch.testing.assert_close(out, _swiglu_reference(x + bias), atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize( + "fp8_input_store, cpu_offload_input", + [(True, False), (False, True)], +) +def test_patched_swiglu_symbol_falls_back_for_unsupported_flags( + fake_megatron_swiglu, fp8_input_store, cpu_offload_input +): + import torch + + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + fused_swiglu, consumers = fake_megatron_swiglu + original = fused_swiglu.bias_swiglu_impl + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + + x = torch.randn(8, 128) + consumers["mlp"].bias_swiglu_impl(x, None, fp8_input_store, cpu_offload_input) + + assert original.call_count == 1 + + +# --------------------------------------------------------------------------- +# 6.5 Cross-kernel isolation. +# --------------------------------------------------------------------------- + + +def test_swiglu_patch_does_not_touch_other_kernels(fake_megatron_swiglu): + """``swiglu=True`` alone must not drag in RMSNorm or CE patching.""" + fused_swiglu, _ = fake_megatron_swiglu + fused_ce, unfused_ce = _install_fake_megatron_ce(tp_size=1) + try: + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + ce_before = fused_ce.fused_vocab_parallel_cross_entropy + unfused_before = unfused_ce.vocab_parallel_cross_entropy + + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + + assert fused_ce.fused_vocab_parallel_cross_entropy is ce_before + assert unfused_ce.vocab_parallel_cross_entropy is unfused_before + assert fused_swiglu.bias_swiglu_impl.__name__ == "liger_bias_swiglu_impl" + finally: + _uninstall_fake_megatron() diff --git a/test/megatron/test_swiglu.py b/test/megatron/test_swiglu.py new file mode 100644 index 000000000..8a500abba --- /dev/null +++ b/test/megatron/test_swiglu.py @@ -0,0 +1,642 @@ +"""Unit tests for LigerMegatronSwiGLU. + +Sections 1-5 deliberately do not import megatron-core; the wrapper's contract is verified +against a local reproduction of Megatron's own reference:: + + y_1, y_2 = torch.chunk(y, 2, -1) + return F.silu(y_1) * y_2 + +The parametrization style mirrors ``test/megatron/test_rms_norm.py``. + +Section 6 is different, and deliberately so. Verifying only against a reimplementation +leaves one class of bug completely uncovered: anything where our *belief* about Megatron's +API is wrong. A signature change, a renamed module path, or a different dispatch condition +would leave every test above green while the integration silently breaks at runtime. Those +tests import the real megatron-core and skip cleanly when it is absent, so they cost +nothing in the default CI environment. Run them with:: + + pip install megatron-core + pytest test/megatron/test_swiglu.py -v +""" + +import os + +import pytest +import torch +import torch.nn.functional as F + +from liger_kernel.megatron.swiglu import LigerMegatronSwiGLU +from liger_kernel.ops.swiglu import LigerFusedGateUpSiLUMulFunction +from liger_kernel.ops.swiglu import LigerSiLUMulFunction +from liger_kernel.utils import infer_device +from test.utils import assert_verbose_allclose +from test.utils import set_seed +from test.utils import supports_bfloat16 + +device = infer_device() + +set_seed(42) +torch.use_deterministic_algorithms(True) + +if device == "cuda": + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + + +# --------------------------------------------------------------------------- +# References + helpers +# --------------------------------------------------------------------------- + + +def _megatron_swiglu_reference(y): + """Byte-for-byte reproduction of megatron.core.fusions.fused_bias_swiglu.swiglu.""" + y_1, y_2 = torch.chunk(y, 2, -1) + return F.silu(y_1) * y_2 + + +def _megatron_bias_swiglu_reference(y, bias): + """Reproduction of Megatron's ``bias_swiglu`` — bias is added before the activation.""" + return _megatron_swiglu_reference(y + bias) + + +class _RecordingFallback: + """Stand-in for Megatron's native ``bias_swiglu_impl``. + + Records every call so tests can assert *that* the fallback was taken and *with what*, + and returns Megatron's reference result so numerical assertions still hold. + """ + + def __init__(self): + self.calls = [] + + def __call__(self, input, bias, fp8_input_store=False, cpu_offload_input=False): + self.calls.append( + { + "input": input, + "bias": bias, + "fp8_input_store": fp8_input_store, + "cpu_offload_input": cpu_offload_input, + } + ) + if bias is not None: + return _megatron_bias_swiglu_reference(input, bias) + return _megatron_swiglu_reference(input) + + +# --------------------------------------------------------------------------- +# Forward + backward correctness +# --------------------------------------------------------------------------- + + +@pytest.mark.flaky(reruns=3, reruns_delay=2) +@pytest.mark.parametrize( + "bs, sl, ffn", + [ + (2, 128, 512), + # Llama-7B-class intermediate size — the shape the kernel is tuned for. + (2, 64, 11008), + # weird shapes + (5, 123, 123), + # single row: exercises the degenerate n_rows == 1 grid. + (1, 1, 64), + ], +) +@pytest.mark.parametrize( + "dtype, atol, rtol", + [ + (torch.float32, 1e-5, 1e-5), + pytest.param( + torch.bfloat16, + 1e-2, + 1e-2, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + ), + ], +) +def test_correctness_3d(bs, sl, ffn, dtype, atol, rtol): + """3D ``[seq, batch, 2*ffn]`` input — the shape Megatron's MLP actually passes.""" + _tensor = torch.randn(bs, sl, 2 * ffn, device=device, dtype=dtype) + do = torch.randn(bs, sl, ffn, device=device, dtype=dtype) + + h1 = _tensor.clone().requires_grad_(True) + h2 = _tensor.clone().requires_grad_(True) + + ref_o = _megatron_swiglu_reference(h1) + ref_o.backward(do, retain_graph=True) + + liger_o = LigerMegatronSwiGLU()(h2) + liger_o.backward(do, retain_graph=True) + + assert liger_o.shape == ref_o.shape + assert_verbose_allclose(liger_o, ref_o, atol=atol, rtol=rtol) + assert_verbose_allclose(h2.grad, h1.grad, atol=atol, rtol=rtol) + + +@pytest.mark.flaky(reruns=3, reruns_delay=2) +@pytest.mark.parametrize("tokens, ffn", [(256, 512), (128, 11008)]) +@pytest.mark.parametrize( + "dtype, atol, rtol", + [ + (torch.float32, 1e-5, 1e-5), + pytest.param( + torch.bfloat16, + 1e-2, + 1e-2, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + ), + ], +) +def test_correctness_2d(tokens, ffn, dtype, atol, rtol): + """2D ``[tokens, 2*ffn]`` input — Megatron's impl accepts both ranks.""" + _tensor = torch.randn(tokens, 2 * ffn, device=device, dtype=dtype) + do = torch.randn(tokens, ffn, device=device, dtype=dtype) + + h1 = _tensor.clone().requires_grad_(True) + h2 = _tensor.clone().requires_grad_(True) + + ref_o = _megatron_swiglu_reference(h1) + ref_o.backward(do, retain_graph=True) + + liger_o = LigerMegatronSwiGLU()(h2) + liger_o.backward(do, retain_graph=True) + + assert liger_o.shape == ref_o.shape + assert_verbose_allclose(liger_o, ref_o, atol=atol, rtol=rtol) + assert_verbose_allclose(h2.grad, h1.grad, atol=atol, rtol=rtol) + + +def test_input_is_not_mutated(): + """The fused gate-up backward must write to a fresh buffer, never into the saved input. + + Megatron keeps a reference to the fc1 output for activation recompute and CUDA-graph + capture. The two-tensor ``LigerSiLUMulFunction`` writes gradients in place into the + buffers it saved; the fused variant deliberately does not, mirroring + ``LigerMegatronRMSNorm``'s ``in_place=False``. + """ + x = torch.randn(4, 8, 64, device=device, dtype=torch.float32) + h = x.clone().requires_grad_(True) + before = h.detach().clone() + + out = LigerMegatronSwiGLU()(h) + out.backward(torch.randn_like(out)) + + assert torch.equal(h.detach(), before), "LigerMegatronSwiGLU mutated its input tensor" + + +@pytest.mark.parametrize( + "dtype, atol, rtol", + [ + (torch.float32, 0.0, 0.0), + pytest.param( + torch.bfloat16, + 0.0, + 0.0, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + ), + ], +) +def test_fused_gate_up_matches_two_tensor_kernel(dtype, atol, rtol): + """The fused kernel must be bit-identical to chunk + ``LigerSiLUMulFunction``. + + The fused path exists purely to remove the copies ``torch.chunk`` forces; it changes + no arithmetic. Asserting exact equality (not allclose) pins that down, so a future + edit to either kernel that silently diverges fails here. + """ + tokens, ffn = 128, 512 + _tensor = torch.randn(tokens, 2 * ffn, device=device, dtype=dtype) + do = torch.randn(tokens, ffn, device=device, dtype=dtype) + + h_fused = _tensor.clone().requires_grad_(True) + fused_o = LigerFusedGateUpSiLUMulFunction.apply(h_fused) + fused_o.backward(do) + + h_split = _tensor.clone().requires_grad_(True) + gate, up = torch.chunk(h_split, 2, dim=-1) + split_o = LigerSiLUMulFunction.apply(gate.contiguous(), up.contiguous()) + split_o.backward(do) + + assert torch.equal(fused_o, split_o) + assert torch.equal(h_fused.grad, h_split.grad) + + +def test_in_place_backward_matches_out_of_place(): + """``in_place=True`` must produce identical gradients, only reusing the input buffer.""" + _tensor = torch.randn(64, 2 * 256, device=device, dtype=torch.float32) + do = torch.randn(64, 256, device=device, dtype=torch.float32) + + h_out = _tensor.clone().requires_grad_(True) + LigerMegatronSwiGLU()(h_out).backward(do) + + h_in = _tensor.clone().requires_grad_(True) + LigerMegatronSwiGLU(in_place=True)(h_in).backward(do) + + assert torch.equal(h_in.grad, h_out.grad) + + +def test_in_place_backward_clobbers_the_saved_input(): + """Pin the documented cost of ``in_place=True``: the fc1 output is destroyed. + + This is not incidental -- it is the whole reason the flag is opt-in. If a future change + makes it non-destructive the flag has become free and the default should be revisited. + """ + x = torch.randn(8, 2 * 64, device=device, dtype=torch.float32) + h = x.clone().requires_grad_(True) + before = h.detach().clone() + + LigerMegatronSwiGLU(in_place=True)(h).backward(torch.randn(8, 64, device=device)) + + assert not torch.equal(h.detach(), before) + + +def test_in_place_rejects_a_second_backward(): + """A repeated backward under ``in_place=True`` must raise, not return wrong gradients. + + Autograd cannot catch this on its own: its version counter is bumped by PyTorch + in-place *ops*, and a Triton kernel writing through the data pointer does not touch it. + Without an explicit guard the second backward recomputes silu from a buffer that now + holds gradients and returns plausible-looking garbage (measured max abs error ~6.8 on a + unit-normal input) with no error at all. This test pins the guard. + """ + x = torch.randn(8, 2 * 64, device=device, dtype=torch.float32) + do = torch.randn(8, 64, device=device) + + h = x.clone().requires_grad_(True) + out = LigerMegatronSwiGLU(in_place=True)(h) + out.backward(do, retain_graph=True) + + with pytest.raises(RuntimeError, match="in_place=True"): + out.backward(do, retain_graph=True) + + +def test_default_path_supports_repeated_backward(): + """The out-of-place default must remain safe to backpropagate through twice.""" + x = torch.randn(8, 2 * 64, device=device, dtype=torch.float32) + do = torch.randn(8, 64, device=device) + + h = x.clone().requires_grad_(True) + out = LigerMegatronSwiGLU()(h) + + out.backward(do, retain_graph=True) + first = h.grad.clone() + h.grad = None + out.backward(do, retain_graph=True) + + assert torch.equal(h.grad, first) + + +# --------------------------------------------------------------------------- +# Fallback behavior for configurations Liger cannot serve +# --------------------------------------------------------------------------- + + +def test_bias_falls_back_to_native(): + fallback = _RecordingFallback() + module = LigerMegatronSwiGLU(fallback_impl=fallback) + + x = torch.randn(16, 64, device=device, dtype=torch.float32) + bias = torch.randn(64, device=device, dtype=torch.float32) + + out = module(x, bias) + + assert len(fallback.calls) == 1 + assert fallback.calls[0]["bias"] is bias + assert_verbose_allclose(out, _megatron_bias_swiglu_reference(x, bias), atol=1e-5, rtol=1e-5) + + +def test_fp8_input_store_falls_back_to_native(): + fallback = _RecordingFallback() + module = LigerMegatronSwiGLU(fallback_impl=fallback) + + x = torch.randn(16, 64, device=device, dtype=torch.float32) + module(x, None, True, False) + + assert len(fallback.calls) == 1 + assert fallback.calls[0]["fp8_input_store"] is True + + +def test_cpu_offload_input_falls_back_to_native(): + fallback = _RecordingFallback() + module = LigerMegatronSwiGLU(fallback_impl=fallback) + + x = torch.randn(16, 64, device=device, dtype=torch.float32) + module(x, None, False, True) + + assert len(fallback.calls) == 1 + assert fallback.calls[0]["cpu_offload_input"] is True + + +def test_supported_config_does_not_touch_fallback(): + """The common path — no bias, no FP8, no offload — must run Liger, not the fallback.""" + fallback = _RecordingFallback() + module = LigerMegatronSwiGLU(fallback_impl=fallback) + + x = torch.randn(16, 64, device=device, dtype=torch.float32) + out = module(x, None, False, False) + + assert fallback.calls == [] + assert_verbose_allclose(out, _megatron_swiglu_reference(x), atol=1e-5, rtol=1e-5) + + +def test_fallback_is_logged_once_per_reason(caplog): + """A fallback taken every microbatch must not spam the log once per step.""" + module = LigerMegatronSwiGLU(fallback_impl=_RecordingFallback()) + x = torch.randn(16, 64, device=device, dtype=torch.float32) + bias = torch.randn(64, device=device, dtype=torch.float32) + + with caplog.at_level("INFO", logger="liger_kernel.megatron.swiglu"): + for _ in range(5): + module(x, bias) + + fallback_records = [r for r in caplog.records if "falling back" in r.getMessage()] + assert len(fallback_records) == 1 + + +def test_unsupported_config_without_fallback_raises_actionable_error(): + """Mode 2 users who never wired a fallback must get a clear error, not silent + numerical drift.""" + module = LigerMegatronSwiGLU() # no fallback_impl + x = torch.randn(16, 64, device=device, dtype=torch.float32) + bias = torch.randn(64, device=device, dtype=torch.float32) + + with pytest.raises(RuntimeError, match="apply_liger_kernel_to_megatron"): + module(x, bias) + + +# --------------------------------------------------------------------------- +# Input validation +# --------------------------------------------------------------------------- + + +def test_rejects_4d_input(): + module = LigerMegatronSwiGLU() + x = torch.randn(2, 2, 4, 64, device=device, dtype=torch.float32) + + with pytest.raises(ValueError, match="2D"): + module(x) + + +def test_rejects_odd_last_dimension(): + """An odd trailing dim means the tensor isn't a concatenated gate/up pair.""" + module = LigerMegatronSwiGLU() + x = torch.randn(16, 63, device=device, dtype=torch.float32) + + with pytest.raises(ValueError, match="even"): + module(x) + + +def test_signature_matches_megatron_positional_order(): + """Megatron's MLP and SharedExpertMLP both call bias_swiglu_impl positionally, so + argument order is part of the contract, not just argument names.""" + import inspect + + params = list(inspect.signature(LigerMegatronSwiGLU.forward).parameters) + assert params == ["self", "input", "bias", "fp8_input_store", "cpu_offload_input"] + + +# --------------------------------------------------------------------------- +# 6. End-to-end validation against a real megatron-core install +# --------------------------------------------------------------------------- +# +# Everything above verifies the wrapper against our own reproduction of Megatron's +# formula. That cannot catch a wrong belief about Megatron's API. The tests below import +# the real package and pin: the true signature and positional order of +# ``bias_swiglu_impl``, numerical parity with Megatron's own TorchScript implementation, +# that the consumer module paths the monkey patch rebinds actually exist and are rebound, +# and that a real ``MLP`` produces identical output and gradients when patched. +# +# The import is conditional rather than a module-level ``pytest.importorskip`` so that the +# dependency-free tests above still run when megatron-core is absent. + +try: + import megatron.core # noqa: F401 + + _MEGATRON_AVAILABLE = True +except ImportError: + _MEGATRON_AVAILABLE = False + +requires_megatron = pytest.mark.skipif( + not _MEGATRON_AVAILABLE + or not (torch.cuda.is_available() or (getattr(torch, "xpu", None) and torch.xpu.is_available())), + reason="requires megatron-core and an accelerator", +) + +_PATCH_TARGETS = ( + "megatron.core.fusions.fused_bias_swiglu", + "megatron.core.transformer.mlp", + "megatron.core.transformer.moe.shared_experts", +) + + +@pytest.fixture +def restore_patch(): + """Snapshot and restore ``bias_swiglu_impl`` in every module the patch touches. + + Restoring only the consumer modules is not enough: the patch is idempotent on the + *defining* module, so leaving that one patched makes a later ``apply_...`` call + short-circuit and the next test observes a half-applied patch. + """ + import importlib + + mods = [importlib.import_module(name) for name in _PATCH_TARGETS] + saved = [(m, m.bias_swiglu_impl) for m in mods] + try: + yield {m.__name__: fn for m, fn in saved} + finally: + for m, fn in saved: + m.bias_swiglu_impl = fn + + +@requires_megatron +def test_real_bias_swiglu_impl_signature_is_what_the_wrapper_assumes(): + """``LigerMegatronSwiGLU.forward`` mirrors this signature positionally. + + Both ``MLP.forward`` and ``SharedExpertMLP.forward`` call ``bias_swiglu_impl`` + positionally, so parameter *order* is part of the contract, not just parameter names. + If Megatron reorders or inserts an argument this fails loudly instead of the patch + silently passing ``fp8_input_store`` where ``bias`` was expected. + """ + import inspect + + from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl + + real = list(inspect.signature(bias_swiglu_impl).parameters) + ours = [p for p in inspect.signature(LigerMegatronSwiGLU.forward).parameters if p != "self"] + + assert real[: len(ours)] == ours, ( + f"megatron-core's bias_swiglu_impl signature {real} no longer matches the order " + f"LigerMegatronSwiGLU.forward accepts {ours}." + ) + + +@requires_megatron +@pytest.mark.parametrize("dtype, atol, rtol", [(torch.float32, 1e-6, 1e-6), (torch.bfloat16, 1e-2, 1e-2)]) +@pytest.mark.parametrize("shape", [(4, 2, 512), (2, 1, 11008)]) +def test_matches_real_megatron_numerically(shape, dtype, atol, rtol): + """Forward and backward parity against Megatron's actual TorchScript implementation.""" + from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl + + s, b, ffn = shape + _tensor = torch.randn(s, b, 2 * ffn, device=device, dtype=dtype) + do = torch.randn(s, b, ffn, device=device, dtype=dtype) + + h_ref = _tensor.clone().requires_grad_(True) + ref = bias_swiglu_impl(h_ref, None, False, False) + ref.backward(do) + + h_liger = _tensor.clone().requires_grad_(True) + got = LigerMegatronSwiGLU()(h_liger, None, False, False) + got.backward(do) + + assert got.shape == ref.shape + torch.testing.assert_close(got, ref, atol=atol, rtol=rtol) + torch.testing.assert_close(h_liger.grad, h_ref.grad, atol=atol, rtol=rtol) + + +@requires_megatron +def test_patch_rebinds_the_real_consumer_modules(restore_patch): + """The by-name imports in the real megatron-core must actually be rebound. + + ``megatron.core.transformer.mlp`` does ``from ...fused_bias_swiglu import + bias_swiglu_impl`` at import time, so patching only the defining module leaves the + consumer holding the original function object. This verifies both that those module + paths still exist and that the rebinding reaches them. + """ + import importlib + + from liger_kernel.megatron.monkey_patch import apply_liger_kernel_to_megatron + + apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=False, swiglu=True) + + for name in _PATCH_TARGETS: + mod = importlib.import_module(name) + patched = mod.bias_swiglu_impl + original = restore_patch[name] + assert patched is not original, f"{name}.bias_swiglu_impl was not rebound" + # The patch installs a plain function wrapping a shared LigerMegatronSwiGLU and + # tags it so re-application is idempotent. __wrapped__ must point back at + # Megatron's real function so unsupported configs fall back rather than recurse. + assert getattr(patched, "__liger_patched__", False), f"{name} rebound to an untagged object" + assert patched.__wrapped__ is original + + # And it must actually compute Liger's result, not just look patched. + mlp_mod = importlib.import_module("megatron.core.transformer.mlp") + x = torch.randn(4, 2, 128, device=device, dtype=torch.float32) + torch.testing.assert_close( + mlp_mod.bias_swiglu_impl(x, None, False, False), + LigerMegatronSwiGLU()(x, None, False, False), + atol=0, + rtol=0, + ) + + +@requires_megatron +def test_patch_rebinds_consumers_imported_after_the_first_call(restore_patch): + """Re-applying the patch must still reach a consumer that reverted to the original. + + Real-world shape of this: a consumer module is imported lazily, *after* + ``apply_liger_kernel_to_megatron`` already ran, so it captures Megatron's original at + its own import time. A naive "already patched, return" short-circuit would leave that + consumer permanently unpatched. + """ + import importlib + + from liger_kernel.megatron.monkey_patch import apply_liger_kernel_to_megatron + + apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=False, swiglu=True) + + mlp_mod = importlib.import_module("megatron.core.transformer.mlp") + mlp_mod.bias_swiglu_impl = restore_patch["megatron.core.transformer.mlp"] + + apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=False, swiglu=True) + + assert getattr(mlp_mod.bias_swiglu_impl, "__liger_patched__", False), ( + "a consumer that reverted to Megatron's original was not re-bound on re-apply" + ) + + +@pytest.fixture(scope="module") +def model_parallel(): + """Minimal single-rank model-parallel init so real Megatron layers can be built.""" + from megatron.core import parallel_state + from megatron.core import tensor_parallel + + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29591") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group(backend="nccl", world_size=1, rank=0) + parallel_state.initialize_model_parallel(tensor_model_parallel_size=1) + # ColumnParallelLinear initializes its weight under the model-parallel RNG fork, which + # does not exist until this is called. Without it, building any parallel layer raises + # "cuda rng state model-parallel-rng is not added". + tensor_parallel.model_parallel_cuda_manual_seed(0) + torch.manual_seed(0) + yield + parallel_state.destroy_model_parallel() + + +def _build_mlp(hidden, ffn, dtype): + from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec + from megatron.core.transformer.mlp import MLP + from megatron.core.transformer.transformer_config import TransformerConfig + + config = TransformerConfig( + num_layers=1, + hidden_size=hidden, + ffn_hidden_size=ffn, + num_attention_heads=4, + gated_linear_unit=True, + activation_func=F.silu, + bias_activation_fusion=True, + add_bias_linear=False, + bf16=(dtype is torch.bfloat16), + params_dtype=dtype, + ) + spec = get_gpt_layer_local_spec().submodules.mlp.submodules + return MLP(config, spec).to(device=device, dtype=dtype) + + +@requires_megatron +def test_real_mlp_forward_backward_parity(model_parallel, restore_patch): + """A real Megatron ``MLP`` must produce identical results with the patch applied. + + This is the test that exercises the actual dispatch path -- ``MLP.forward`` deciding to + call ``bias_swiglu_impl`` based on ``config.bias_activation_fusion`` / + ``gated_linear_unit`` / ``activation_func``, with the real ``linear_fc1`` output layout + feeding it. + + One MLP instance is reused for both runs rather than building two and copying weights: + the patch is a module-level symbol swap, so the same instance picks it up, and this + sidesteps Megatron's ``_extra_state`` entries (which are None and break a naive + ``state_dict`` round-trip). + """ + import importlib + + from liger_kernel.megatron.monkey_patch import apply_liger_kernel_to_megatron + + hidden, ffn, dtype = 128, 512, torch.float32 + mlp_mod = importlib.import_module("megatron.core.transformer.mlp") + original = restore_patch["megatron.core.transformer.mlp"] + + mlp = _build_mlp(hidden, ffn, dtype) + x = torch.randn(8, 2, hidden, device=device, dtype=dtype) + + def run(): + for p_ in mlp.parameters(): + p_.grad = None + h = x.clone().requires_grad_(True) + out, _ = mlp(h) + out.sum().backward() + return out.detach().clone(), h.grad.clone(), [p_.grad.clone() for p_ in mlp.parameters()] + + ref_out, ref_dx, ref_dw = run() + + apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=False, swiglu=True) + assert mlp_mod.bias_swiglu_impl is not original, "patch did not take effect" + liger_out, liger_dx, liger_dw = run() + + torch.testing.assert_close(liger_out, ref_out, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(liger_dx, ref_dx, atol=1e-5, rtol=1e-5) + for got, want in zip(liger_dw, ref_dw): + torch.testing.assert_close(got, want, atol=1e-5, rtol=1e-5) From 9bfdbffa19da4b57e1a25e6a8d5108ed643cc3e4 Mon Sep 17 00:00:00 2001 From: Buffett Liu Date: Tue, 28 Jul 2026 23:27:45 -0700 Subject: [PATCH 2/2] Fix and revision --- examples/megatron/README.md | 52 +-- examples/megatron/run_mode1_monkey_patch.py | 44 +-- examples/megatron/run_mode2_hand_spec.py | 18 +- src/liger_kernel/megatron/__init__.py | 9 +- src/liger_kernel/megatron/monkey_patch.py | 168 ++++----- src/liger_kernel/megatron/swiglu.py | 14 +- src/liger_kernel/ops/swiglu.py | 58 +-- test/megatron/test_monkey_patch.py | 380 +++++++++++--------- test/megatron/test_swiglu.py | 262 +------------- 9 files changed, 359 insertions(+), 646 deletions(-) diff --git a/examples/megatron/README.md b/examples/megatron/README.md index 126dc7888..abf06e950 100644 --- a/examples/megatron/README.md +++ b/examples/megatron/README.md @@ -2,29 +2,39 @@ Two self-contained scripts demonstrating the integration modes shipped by `liger_kernel.megatron`. Both train a tiny GPT model with mock data on -2 × GPU (TP=1, PP=1, DP=2) for 5 iterations and print the resolved norm -classes, CE bindings and SwiGLU bindings so you can see which slots picked -up Liger. +2 × GPU (TP=1, PP=1, DP=2) for 5 iterations and print the resolved bindings +so you can see which slots picked up Liger. + +## Op support matrix + +| Op | Mode 1 flag | Patched symbol(s) | Mode 2 class | Slot used in Mode 2 | +|---|---|---|---|---| +| RMSNorm | `rms_norm=True` (on by default) | `LocalSpecProvider.layer_norm`, `transformer_block.LayerNormImpl` | `LigerMegatronRMSNorm` | every norm slot, incl. block-level `final_layernorm` | +| Cross-entropy | `cross_entropy=True` (opt-in) | `fused_cross_entropy.fused_vocab_parallel_cross_entropy`, `tensor_parallel.cross_entropy.vocab_parallel_cross_entropy` | `LigerMegatronCrossEntropy` | none — `GPTModel` subclass overriding `compute_language_model_loss` | +| SwiGLU | `swiglu=True` (opt-in) | `fusions.fused_bias_swiglu.SwiGLUFunction` | `LigerMegatronSwiGLU` | the `mlp` module slot — an `MLP` subclass | + +Notes that apply to the table: + +- RMSNorm only covers the local (non-TE) backend. +- SwiGLU is used only when `gated_linear_unit=True`, `activation_func=F.silu`, + and `bias_activation_fusion=True`; otherwise the patch is applied but not + exercised. +- Liger replaces only `SwiGLUFunction`, so bias and MoE variants stay on + Megatron. +- Cross-entropy and SwiGLU are wired through subclasses (no dedicated spec + slot). ## Prerequisites - A working Megatron-Core install (`pip install megatron-core`). - `liger-kernel` installed (editable or from PyPI). -- `psutil` (used by Megatron's async checkpoint worker pool). +- `psutil`. - At least 2 GPUs. ## Mode 1 — `apply_liger_kernel_to_megatron()` (monkey-patch) -One-line opt-in. Patches `LocalSpecProvider.layer_norm` and -`transformer_block.LayerNormImpl` so every RMSNorm slot becomes Liger, the -two cross-entropy entry points, and `bias_swiglu_impl` — all without -changing the spec the user constructs. - -Note the SwiGLU dispatch flags on `TransformerConfig` -(`gated_linear_unit=True`, `activation_func=F.silu`, -`bias_activation_fusion=True`). `MLP.forward` only routes through -`bias_swiglu_impl` when all three hold, so without them the patch is -applied but never reached — no speedup and no error. +One-line opt-in. Patches the symbols in the matrix above without changing +the spec the user constructs. ```bash torchrun --nproc_per_node=2 \ @@ -34,17 +44,9 @@ torchrun --nproc_per_node=2 \ ## Mode 2 — hand-assembled `TransformerBlockSubmodules` -Slot-level control. Explicitly places `LigerMegatronRMSNorm` into each -norm slot, including the block-level `final_layernorm`, and a -`LigerMegatronSwiGLU`-backed `MLP` subclass into the `mlp` slot. Useful -when you want to mix Liger with other backends (e.g. TransformerEngine) on -a per-slot basis. - -SwiGLU has no spec slot of its own — `MLPSubmodules.activation_func` is -only read when `config.use_te_activation_func` is set, so it is the -TransformerEngine hook rather than a general one. The MLP *class* is a -slot, so this script subclasses `MLP`, the same way it subclasses -`GPTModel` for cross-entropy. +Slot-level control — see the "Mode 2 class" column of the matrix above. +Useful when you want to mix Liger with other backends (e.g. +TransformerEngine) on a per-slot basis. ```bash torchrun --nproc_per_node=2 \ diff --git a/examples/megatron/run_mode1_monkey_patch.py b/examples/megatron/run_mode1_monkey_patch.py index 64563d493..9e5467e9c 100644 --- a/examples/megatron/run_mode1_monkey_patch.py +++ b/examples/megatron/run_mode1_monkey_patch.py @@ -11,9 +11,8 @@ (the fused CE path) - ``tensor_parallel.cross_entropy.vocab_parallel_cross_entropy`` (the unfused CE path) - - ``fusions.fused_bias_swiglu.bias_swiglu_impl`` *and* the by-name - re-exports of it in ``transformer.mlp`` and - ``transformer.moe.shared_experts`` (the SwiGLU activation) + - ``fusions.fused_bias_swiglu.SwiGLUFunction`` (the SwiGLU + activation, reached via Megatron's own ``bias_swiglu_impl``) 2. ``normalization="RMSNorm"`` added to ``TransformerConfig`` so the model actually has RMSNorm slots to patch (Megatron defaults to @@ -109,23 +108,9 @@ def model_provider() -> GPTModel: ) -def _ensure_dataset_helpers() -> None: - """Build Megatron's C++ dataset helpers only if not already compiled. - - ``compile_helpers()`` shells out to ``make``, which fails on pip-installed - megatron-core wheels (no Makefile, but the extension is already present). - """ - try: - from megatron.core.datasets import helpers_cpp # noqa: F401 - - return - except ImportError: - compile_helpers() - - def get_train_data_iterator() -> Iterator: if torch.distributed.get_rank() == 0: - _ensure_dataset_helpers() + compile_helpers() torch.distributed.barrier() cfg = GPTDatasetConfig( random_seed=0, @@ -134,7 +119,7 @@ def get_train_data_iterator() -> Iterator: reset_attention_mask=False, eod_mask_loss=False, tokenizer=MegatronTokenizer.from_pretrained( - metadata_path={"library": "null"}, + metadata_path={"library": "null-text"}, vocab_size=_SEQUENCE_LENGTH, ), mid_level_dataset_surplus=0.005, @@ -188,24 +173,27 @@ def _print_ce_symbols() -> None: def _print_swiglu_symbols() -> None: - """Show ``bias_swiglu_impl`` bindings in the defining module and its consumers. + """Show which ``SwiGLUFunction`` the dense MLP and shared experts end up running. - The two consumer modules capture the symbol at import time, so the defining - module and each consumer are printed separately to verify all were rebound. + Liger replaces the class, not ``bias_swiglu_impl``. That function is a plain + dispatcher that looks the class up in its own module globals on every call, + so the two consumer modules keep their import-time ``bias_swiglu_impl`` + binding and still route to Liger. Both facts are printed below. """ import megatron.core.fusions.fused_bias_swiglu as defining import megatron.core.transformer.mlp as mlp import megatron.core.transformer.moe.shared_experts as shared + tag = "Liger" if getattr(defining.SwiGLUFunction, "__liger_patched__", False) else "Megatron" print("\n=== Resolved SwiGLU symbols ===") + print(f" {'fusions.fused_bias_swiglu.SwiGLUFunction':52s} \u2192 [{tag}]") for label, mod in ( - ("fusions.fused_bias_swiglu (defines)", defining), - ("transformer.mlp (consumes)", mlp), - ("transformer.moe.shared_experts (consumes)", shared), + ("transformer.mlp", mlp), + ("transformer.moe.shared_experts", shared), ): - fn = mod.bias_swiglu_impl - tag = "Liger" if getattr(fn, "__liger_patched__", False) else "Megatron" - print(f" {label:44s} → {fn.__name__} [{tag}]") + same = mod.bias_swiglu_impl is defining.bias_swiglu_impl + note = "unpatched, resolves the class above" if same else "REBOUND \u2014 unexpected" + print(f" {label + '.bias_swiglu_impl':52s} \u2192 {note}") print() diff --git a/examples/megatron/run_mode2_hand_spec.py b/examples/megatron/run_mode2_hand_spec.py index 268862d64..b24036638 100644 --- a/examples/megatron/run_mode2_hand_spec.py +++ b/examples/megatron/run_mode2_hand_spec.py @@ -219,23 +219,9 @@ def model_provider() -> GPTModel: ) -def _ensure_dataset_helpers() -> None: - """Build Megatron's C++ dataset helpers only if not already compiled. - - ``compile_helpers()`` shells out to ``make``, which fails on pip-installed - megatron-core wheels (no Makefile, but the extension is already present). - """ - try: - from megatron.core.datasets import helpers_cpp # noqa: F401 - - return - except ImportError: - compile_helpers() - - def get_train_data_iterator() -> Iterator: if torch.distributed.get_rank() == 0: - _ensure_dataset_helpers() + compile_helpers() torch.distributed.barrier() cfg = GPTDatasetConfig( random_seed=0, @@ -244,7 +230,7 @@ def get_train_data_iterator() -> Iterator: reset_attention_mask=False, eod_mask_loss=False, tokenizer=MegatronTokenizer.from_pretrained( - metadata_path={"library": "null"}, + metadata_path={"library": "null-text"}, vocab_size=_SEQUENCE_LENGTH, ), mid_level_dataset_surplus=0.005, diff --git a/src/liger_kernel/megatron/__init__.py b/src/liger_kernel/megatron/__init__.py index d9a88f148..94f3da552 100644 --- a/src/liger_kernel/megatron/__init__.py +++ b/src/liger_kernel/megatron/__init__.py @@ -5,10 +5,11 @@ LayerNormBuilder protocol. LigerMegatronCrossEntropy — drop-in for Megatron-LM's vocab-parallel cross-entropy (Megatron defaults). Supports all TP sizes. - LigerMegatronSwiGLU — drop-in for Megatron-Core's ``bias_swiglu_impl``, - the fused gated-SiLU activation used by the dense MLP and the MoE - shared experts. Falls back to Megatron's implementation for bias, - FP8 input store, and CPU activation offload. + LigerMegatronSwiGLU — MLP subclass for Mode 2, replacing the fused + gated-SiLU activation used by the dense MLP and the MoE shared + experts. Mode 1 patches ``fused_bias_swiglu.SwiGLUFunction`` instead; + both fall back to Megatron for FP8 input store and CPU activation + offload, and neither touches the bias or MoE-routed variants. apply_liger_kernel_to_megatron — patches Megatron-Core so existing training scripts pick up Liger kernels with one line. Currently supports RMSNorm (via BackendSpecProvider), both the fused and unfused diff --git a/src/liger_kernel/megatron/monkey_patch.py b/src/liger_kernel/megatron/monkey_patch.py index fd8de5f7d..e97567ca3 100644 --- a/src/liger_kernel/megatron/monkey_patch.py +++ b/src/liger_kernel/megatron/monkey_patch.py @@ -2,7 +2,6 @@ from __future__ import annotations -import inspect import logging logger = logging.getLogger(__name__) @@ -41,14 +40,11 @@ def apply_liger_kernel_to_megatron( argument, matching native's ``(logits, target, label_smoothing=0.0, tp_group=None)``. swiglu: When ``True`` replace - ``megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl`` with Liger's - Triton SiLU-multiply kernel. Covers the dense ``MLP`` and the MoE - ``SharedExpertMLP``, i.e. every site Megatron reaches when - ``config.bias_activation_fusion=True``, ``config.gated_linear_unit=True`` - and ``config.activation_func is F.silu``. Default ``False`` so adopters - opt in explicitly. Calls Liger cannot serve (a non-``None`` bias, - ``activation_func_fp8_input_store=True``, or CPU activation offloading) - transparently defer to Megatron's original implementation. + ``megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction`` with Liger's + Triton SiLU-multiply kernel, covering the dense ``MLP`` and the MoE + ``SharedExpertMLP``. Default ``False`` so adopters opt in explicitly. + See ``_patch_swiglu_function`` for the configurations this reaches + and the ones that stay on Megatron. Notes: Call this BEFORE building your model. Patching after instantiation @@ -73,7 +69,7 @@ def apply_liger_kernel_to_megatron( _patch_fused_vocab_parallel_cross_entropy() _patch_vocab_parallel_cross_entropy() if swiglu: - _patch_bias_swiglu_impl() + _patch_swiglu_function() def _patch_local_spec_provider_layer_norm() -> None: @@ -85,20 +81,11 @@ def _patch_local_spec_provider_layer_norm() -> None: return # already patched original_layer_norm = backends.LocalSpecProvider.layer_norm - _original_sig = inspect.signature(original_layer_norm) - - def patched_layer_norm(self, *args, **kwargs): - # Forward through the *original* signature rather than restating it. Megatron has - # changed this method's parameters across releases (``has_residual`` was added - # after ``for_qk``), and a hardcoded forwarding call raises TypeError on any - # version whose signature does not match exactly. Binding to the real signature - # also means ``rms_norm`` is read correctly whether the caller passed it - # positionally or by keyword. - bound = _original_sig.bind(self, *args, **kwargs) - bound.apply_defaults() - if bound.arguments.get("rms_norm", False): + + def patched_layer_norm(self, rms_norm: bool = False, for_qk: bool = False, has_residual: bool = False, **kwargs): + if rms_norm: return LigerMegatronRMSNorm - return original_layer_norm(*bound.args, **bound.kwargs) + return original_layer_norm(self, rms_norm=rms_norm, for_qk=for_qk, has_residual=has_residual, **kwargs) setattr(patched_layer_norm, _PATCH_MARKER, True) setattr(patched_layer_norm, "__wrapped__", original_layer_norm) @@ -279,27 +266,28 @@ def liger_vocab_parallel_cross_entropy( ) -# Modules that do ``from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl`` -# and therefore hold their own binding. Must be rebound individually after patching. -# Verified against Megatron-LM ``main``; missing entries are skipped. -_BIAS_SWIGLU_CONSUMER_MODULES = ( - "megatron.core.transformer.mlp", - "megatron.core.transformer.moe.shared_experts", -) - - -def _patch_bias_swiglu_impl() -> None: - """Replace ``megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl`` with Liger. +def _patch_swiglu_function() -> None: + """Replace ``megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction`` with Liger. Covers ``MLP.forward`` and ``SharedExpertMLP.forward`` when ``config.bias_activation_fusion=True``, ``config.gated_linear_unit=True`` and ``config.activation_func is F.silu``. - Not patched: ``weighted_bias_swiglu_impl`` (needs routing-weight grad), - the unfused ``glu()`` closure (no module-level symbol), and - ``config.use_te_activation_func=True`` (TE owns that path). - - Unsupported configs (bias, FP8, CPU offload) defer to Megatron's original at call time. + Megatron's ``bias_swiglu_impl`` is intentionally left alone: it is a plain function + that resolves ``SwiGLUFunction`` from its own globals on every call, so replacing the + class reaches every caller -- including modules that already did + ``from ... import bias_swiglu_impl`` -- regardless of import order. Nothing outside + ``fused_bias_swiglu`` references ``SwiGLUFunction``, so there is exactly one binding to + replace and no stale copies can exist. + + Patching one layer further down (the ``swiglu`` / ``swiglu_back`` math helpers) is not + an option: they are ``@jit_fuser``-decorated, and ``jit_fuser`` is ``torch.jit.script`` + below torch 2.2, which compiles them at import time -- the patch would silently no-op. + + Not patched: ``BiasSwiGLUFunction`` (Liger's kernel has no bias term, so a non-``None`` + bias keeps using Megatron by construction), ``WeightedSwiGLUFunction`` (MoE routed + experts, needs routing-weight grad), and ``config.use_te_activation_func=True`` + (TransformerEngine owns that path). """ try: import megatron.core.fusions.fused_bias_swiglu as fused_swiglu @@ -307,77 +295,57 @@ def _patch_bias_swiglu_impl() -> None: raise ImportError( "apply_liger_kernel_to_megatron(swiglu=True) requires megatron-core to be " "installed. Expected symbol path: " - "megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl." + "megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction." ) from exc - if not hasattr(fused_swiglu, "bias_swiglu_impl"): + if not hasattr(fused_swiglu, "SwiGLUFunction"): raise ImportError( - "megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl not found. The symbol " + "megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction not found. The symbol " "path may have changed in your Megatron-LM version. Please file an issue on " "https://github.com/linkedin/Liger-Kernel with your megatron-core version." ) - if getattr(fused_swiglu.bias_swiglu_impl, _PATCH_MARKER, False): - # Already patched -- re-run consumer rebinding for modules imported after the - # first patch call (they'd still hold Megatron's original binding). - liger_impl = fused_swiglu.bias_swiglu_impl - _rebind_consumer_symbol( - _BIAS_SWIGLU_CONSUMER_MODULES, - "bias_swiglu_impl", - liger_impl.__wrapped__, - liger_impl, - ) - return - - original = fused_swiglu.bias_swiglu_impl - - from liger_kernel.megatron.swiglu import LigerMegatronSwiGLU - - # Reuse a single instance: stateless except for the fallback-log dedup set. - swiglu_module = LigerMegatronSwiGLU(fallback_impl=original) - - def liger_bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False): - return swiglu_module(input, bias, fp8_input_store, cpu_offload_input) - - setattr(liger_bias_swiglu_impl, _PATCH_MARKER, True) - setattr(liger_bias_swiglu_impl, "__wrapped__", original) - fused_swiglu.bias_swiglu_impl = liger_bias_swiglu_impl + if getattr(fused_swiglu.SwiGLUFunction, _PATCH_MARKER, False): + return # already patched - patched_consumers = _rebind_consumer_symbol( - _BIAS_SWIGLU_CONSUMER_MODULES, - "bias_swiglu_impl", - original, - liger_bias_swiglu_impl, - ) + original = fused_swiglu.SwiGLUFunction - logger.info( - "Patched megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl with Liger SwiGLU (also rebound in: %s).", - ", ".join(patched_consumers) if patched_consumers else "no already-imported consumers", - ) + from liger_kernel.ops.swiglu import LigerFusedGateUpSiLUMulFunction + # Deduplicate fallback logs so unsupported configs do not spam every step. + logged_fallbacks = set() -def _rebind_consumer_symbol( - module_names, - attr_name: str, - original, - replacement, -): - """Point already-imported ``from X import attr`` bindings at ``replacement``. + class _LigerSwiGLUFunction: + """Adapter matching Megatron's ``SwiGLUFunction.apply`` signature. - Only rebinds modules already in ``sys.modules`` whose binding is still ``original``. - Modules imported later pick up the patched symbol naturally. + This is not an ``autograd.Function``. It exists to map Megatron's + ``(input, fp8_input_store, cpu_offload_input)`` call shape to Liger's kernel + signature without changing call sites. + """ - Returns the list of module names actually rebound. - """ - import sys - - rebound = [] - for name in module_names: - module = sys.modules.get(name) - if module is None: - continue # not imported yet - if getattr(module, attr_name, None) is not original: - continue # absent, or already pointing somewhere else - setattr(module, attr_name, replacement) - rebound.append(name) - return rebound + @staticmethod + def apply(input, fp8_input_store=False, cpu_offload_input=False): + # FP8 input-store and CPU offload use Megatron-specific backward storage + # semantics, so these paths defer to native. + reason = None + if fp8_input_store: + reason = "config.activation_func_fp8_input_store=True" + elif cpu_offload_input: + reason = "CPU activation offloading enabled" + if reason is not None: + if reason not in logged_fallbacks: + logged_fallbacks.add(reason) + logger.info( + "Liger SwiGLU is deferring to Megatron's native SwiGLUFunction: %s. " + "Numerics and memory behavior are unchanged for this configuration.", + reason, + ) + return original.apply(input, fp8_input_store, cpu_offload_input) + + return LigerFusedGateUpSiLUMulFunction.apply(input, False) + + setattr(_LigerSwiGLUFunction, _PATCH_MARKER, True) + setattr(_LigerSwiGLUFunction, "__wrapped__", original) + fused_swiglu.SwiGLUFunction = _LigerSwiGLUFunction + + logger.info("Patched megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction with Liger SwiGLU.") diff --git a/src/liger_kernel/megatron/swiglu.py b/src/liger_kernel/megatron/swiglu.py index 7651ef8af..639b4c765 100644 --- a/src/liger_kernel/megatron/swiglu.py +++ b/src/liger_kernel/megatron/swiglu.py @@ -1,7 +1,12 @@ """Megatron-Core compatible SwiGLU backed by the Liger Triton SiLU-multiply kernel. -Drop-in for ``megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl``, which is the -symbol ``MLP.forward`` and ``SharedExpertMLP.forward`` call when +``LigerMegatronSwiGLU`` is an ``nn.Module`` mirroring ``bias_swiglu_impl``'s signature, +for hand-assembled specs (Mode 2) where the caller owns the MLP class. Mode 1 does not use +it: the monkey patch installs a small adapter over +``megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction`` instead -- see +``monkey_patch._patch_swiglu_function``. + +Both cover the path ``MLP.forward`` and ``SharedExpertMLP.forward`` take when ``config.bias_activation_fusion=True``, ``config.gated_linear_unit=True`` and ``config.activation_func is F.silu``. @@ -37,11 +42,6 @@ import torch import torch.nn as nn -# Force-import the submodule so liger_kernel.ops.swiglu can resolve -# torch.distributed.tensor.DTensor on torch 2.11+, where the subpackage is no longer -# auto-loaded as an attribute of torch.distributed. -import torch.distributed.tensor # noqa: F401 # isort: skip - from liger_kernel.ops import LigerFusedGateUpSiLUMulFunction logger = logging.getLogger(__name__) diff --git a/src/liger_kernel/ops/swiglu.py b/src/liger_kernel/ops/swiglu.py index c1cf50d6f..25053b04b 100644 --- a/src/liger_kernel/ops/swiglu.py +++ b/src/liger_kernel/ops/swiglu.py @@ -286,14 +286,14 @@ def backward(ctx, dc): # --------------------------------------------------------------------------- # Fused gate-up variant # --------------------------------------------------------------------------- -# For fused ``[tokens, 2*n]`` gate-up tensors (Megatron, HF ``gate_up_proj``): the -# kernels read both halves via a column offset into the single buffer -- no copies, -# no cat. Input row stride is ``2*n``, output is ``n``. +# For fused ``[tokens, 2 * ffn_size]`` gate-up tensors (Megatron, HF ``gate_up_proj``): +# the kernels read both halves via a column offset into the single buffer -- no copies, +# no cat. Input row stride is ``2 * ffn_size``, output is ``ffn_size``. @triton.jit def _swiglu_fused_gate_up_forward_kernel( - y_ptr, c_ptr, in_stride, out_stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr + y_ptr, c_ptr, in_stride, out_stride, ffn_size: tl.constexpr, BLOCK_SIZE: tl.constexpr ): program_id = tl.program_id(0).to(tl.int64) @@ -301,17 +301,17 @@ def _swiglu_fused_gate_up_forward_kernel( c_ptr += program_id * out_stride col_offsets = tl.arange(0, BLOCK_SIZE) - mask = col_offsets < n_cols + mask = col_offsets < ffn_size - # Gate occupies columns [0, n), up occupies [n, 2n) of the same row. + # Gate occupies columns [0, ffn_size), up occupies [ffn_size, 2 * ffn_size) of the same row. gate = tl.load(y_ptr + col_offsets, mask=mask, other=0).to(tl.float32) - up = tl.load(y_ptr + n_cols + col_offsets, mask=mask, other=0) + up = tl.load(y_ptr + ffn_size + col_offsets, mask=mask, other=0) tl.store(c_ptr + col_offsets, silu(gate).cast(up.dtype) * up, mask=mask) @triton.jit def _swiglu_fused_gate_up_backward_kernel( - dc_ptr, y_ptr, dy_ptr, in_stride, out_stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr + dc_ptr, y_ptr, dy_ptr, in_stride, out_stride, ffn_size: tl.constexpr, BLOCK_SIZE: tl.constexpr ): program_id = tl.program_id(0).to(tl.int64) @@ -320,11 +320,11 @@ def _swiglu_fused_gate_up_backward_kernel( dy_ptr += program_id * in_stride col_offsets = tl.arange(0, BLOCK_SIZE) - mask = col_offsets < n_cols + mask = col_offsets < ffn_size dc = tl.load(dc_ptr + col_offsets, mask=mask, other=0) gate = tl.load(y_ptr + col_offsets, mask=mask, other=0).to(tl.float32) - up = tl.load(y_ptr + n_cols + col_offsets, mask=mask, other=0) + up = tl.load(y_ptr + ffn_size + col_offsets, mask=mask, other=0) # Recompute silu from saved input. When dy_ptr aliases y_ptr (in_place=True), all # loads precede all stores and each program owns one row, so aliasing is safe here. @@ -335,59 +335,59 @@ def _swiglu_fused_gate_up_backward_kernel( d_up = dc * silu_gate tl.store(dy_ptr + col_offsets, d_gate, mask=mask) - tl.store(dy_ptr + n_cols + col_offsets, d_up, mask=mask) + tl.store(dy_ptr + ffn_size + col_offsets, d_up, mask=mask) def swiglu_fused_gate_up_forward(y): - """SwiGLU over a fused ``[..., 2 * n]`` gate-up tensor. Returns ``(y, c)``.""" + """SwiGLU over a fused ``[..., 2 * ffn_size]`` gate-up tensor. Returns ``(y, c)``.""" ori_shape = y.shape - two_n = ori_shape[-1] - if two_n % 2 != 0: - raise ValueError(f"fused gate-up input must have an even trailing dim; got {two_n}.") - n_cols = two_n // 2 + fused_size = ori_shape[-1] + if fused_size % 2 != 0: + raise ValueError(f"fused gate-up input must have an even trailing dim; got {fused_size}.") + ffn_size = fused_size // 2 - y = y.view(-1, two_n) + y = y.view(-1, fused_size) n_rows = y.shape[0] - c = torch.empty(n_rows, n_cols, dtype=y.dtype, device=y.device) + c = torch.empty(n_rows, ffn_size, dtype=y.dtype, device=y.device) - BLOCK_SIZE, num_warps = calculate_settings(n_cols) + BLOCK_SIZE, num_warps = calculate_settings(ffn_size) _swiglu_fused_gate_up_forward_kernel[(n_rows,)]( y, c, y.stride(-2), c.stride(-2), - n_cols=n_cols, + ffn_size=ffn_size, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, ) - return y, c.view(*ori_shape[:-1], n_cols) + return y, c.view(*ori_shape[:-1], ffn_size) def swiglu_fused_gate_up_backward(y, dc, in_place=False): - """Gradient w.r.t. the fused ``[..., 2 * n]`` gate-up tensor. + """Gradient w.r.t. the fused ``[..., 2 * ffn_size]`` gate-up tensor. Args: in_place: Overwrite ``y`` with the gradient in place instead of allocating a new - ``[..., 2*n]`` buffer. Saves ~1 GB peak on H100 at ``[2048*4, 2*32768]`` bf16. + ``[..., 2 * ffn_size]`` buffer. Saves ~1 GB peak on H100 at ``[2048*4, 2*32768]`` bf16. Default False -- Megatron may hold references to the fc1 output for activation recompute or CUDA-graph capture, so clobbering it requires care. """ - two_n = y.shape[-1] - n_cols = two_n // 2 + fused_size = y.shape[-1] + ffn_size = fused_size // 2 - y = y.view(-1, two_n) - dc = dc.view(-1, n_cols) + y = y.view(-1, fused_size) + dc = dc.view(-1, ffn_size) n_rows = dc.shape[0] dy = y if in_place else torch.empty_like(y) - BLOCK_SIZE, num_warps = calculate_settings(n_cols) + BLOCK_SIZE, num_warps = calculate_settings(ffn_size) _swiglu_fused_gate_up_backward_kernel[(n_rows,)]( dc, y, dy, y.stride(-2), dc.stride(-2), - n_cols=n_cols, + ffn_size=ffn_size, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, ) diff --git a/test/megatron/test_monkey_patch.py b/test/megatron/test_monkey_patch.py index 6bef010e2..634a4abd5 100644 --- a/test/megatron/test_monkey_patch.py +++ b/test/megatron/test_monkey_patch.py @@ -60,7 +60,6 @@ def _install_fake_megatron_ce( Returns a tuple ``(fused_ce_module, unfused_ce_module)`` so tests can inspect what the patch helpers wrote onto them. """ - _snapshot_real_megatron() _, megatron_core = _ensure_megatron_roots() fusions = types.ModuleType("megatron.core.fusions") fused_ce = types.ModuleType("megatron.core.fusions.fused_cross_entropy") @@ -121,7 +120,6 @@ def _install_fake_megatron_rms_norm( level patch only displaces that fallback; set False to verify the no-op path taken under TE / Apex. """ - _snapshot_real_megatron() _, megatron_core = _ensure_megatron_roots() backends = None @@ -191,73 +189,31 @@ class _SomeOtherNorm: return backends, transformer_block -_SAVED_REAL_MEGATRON: dict = {} -_STUBS_ACTIVE = False - - -def _is_megatron_module(name: str) -> bool: - return name == "megatron" or name.startswith("megatron.") - - -def _snapshot_real_megatron(): - """Record the entire ``megatron*`` region of ``sys.modules`` before stubbing it. - - Only has an effect when megatron-core is genuinely installed, e.g. an environment set - up to run the end-to-end section of ``test_swiglu.py``. See ``_uninstall_fake_megatron`` - for why this has to cover the whole package rather than the specific names the stubs - replace. - """ - global _STUBS_ACTIVE - if _STUBS_ACTIVE: - return # a stub install is already in effect; don't snapshot the stubs themselves - _SAVED_REAL_MEGATRON.clear() - _SAVED_REAL_MEGATRON.update({name: mod for name, mod in sys.modules.items() if _is_megatron_module(name)}) - _STUBS_ACTIVE = True - - def _uninstall_fake_megatron(): - """Restore ``sys.modules`` to exactly its pre-stub state. - - Two subtleties, both of which only bite when a real megatron-core is installed: - - 1. Popping a real module and walking away is not neutral. The next import rebuilds it - as a *new* object with fresh globals, silently discarding Megatron's process-global - parallel state -- so a later ``initialize_model_parallel`` appears not to have run. - 2. Evicting a hand-picked list of names is not enough. Submodules outside the list - (e.g. ``megatron.core.tensor_parallel.random``) survive holding references to the - evicted parent, so the process ends up with two live ``parallel_state`` objects and - whichever one gets initialized is not the one that gets read. - - Clearing and restoring the whole ``megatron*`` prefix avoids both. - """ - global _STUBS_ACTIVE - if not _STUBS_ACTIVE: - return # nothing was stubbed; leave a real install alone - - stubbed = [n for n in list(sys.modules) if _is_megatron_module(n)] - for name in stubbed: - del sys.modules[name] - sys.modules.update(_SAVED_REAL_MEGATRON) - - # sys.modules is not the only place a stub lands. Importing `megatron.core.x` also - # binds `x` as an *attribute* of the `megatron.core` package object, and restoring - # sys.modules does not undo that binding. A later `from megatron.core import x` reads - # the attribute and gets the stub straight back -- leaving two live copies of a module - # whose globals are process-wide state, which is how `initialize_model_parallel` ends - # up writing to one copy while everything else reads the other. - for name in stubbed: - parent_name, _, child = name.rpartition(".") - parent = sys.modules.get(parent_name) - if parent is None: - continue - real = _SAVED_REAL_MEGATRON.get(name) - if real is not None: - setattr(parent, child, real) - elif hasattr(parent, child): - delattr(parent, child) - - _SAVED_REAL_MEGATRON.clear() - _STUBS_ACTIVE = False + """Tear down every stub module installed by any installer.""" + for mod in [ + # CE side + "megatron.core.parallel_state", + "megatron.core.fusions.fused_cross_entropy", + # RMSNorm side + "megatron.core.models.backends", + "megatron.core.models", + "megatron.core.transformer.transformer_block", + "megatron.core.transformer.torch_norm", + # SwiGLU side + "megatron.core.fusions.fused_bias_swiglu", + "megatron.core.transformer.mlp", + "megatron.core.transformer.moe.shared_experts", + "megatron.core.transformer.moe", + # Shared packages + "megatron.core.fusions", + "megatron.core.tensor_parallel.cross_entropy", + "megatron.core.tensor_parallel", + "megatron.core.transformer", + "megatron.core", + "megatron", + ]: + sys.modules.pop(mod, None) @pytest.fixture @@ -286,37 +242,70 @@ def _install_fake_megatron_swiglu( Returns ``(fused_swiglu_module, consumer_modules_dict)``. - The consumer modules matter here in a way they don't for CE or RMSNorm: real Megatron - does ``from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl`` in - ``transformer/mlp.py`` and ``transformer/moe/shared_experts.py``, so each holds its own - binding. A patch that only rewrites the defining module leaves every call site on - native Megatron. The stub reproduces those by-name imports so the tests can catch that. + The stub mirrors real Megatron's layering, because the layering is what the patch + depends on: ``bias_swiglu_impl`` is a plain function that resolves ``SwiGLUFunction`` + from its own module globals on every call, and Liger replaces the class rather than + the function. The consumer modules reproduce + ``from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl`` so the tests + can prove those stale bindings still reach Liger without being rebound. """ - _snapshot_real_megatron() _, megatron_core = _ensure_megatron_roots() # ``megatron.core.fusions`` may already exist if the CE installer ran first. fusions = sys.modules.get("megatron.core.fusions") or types.ModuleType("megatron.core.fusions") fused_swiglu = types.ModuleType("megatron.core.fusions.fused_bias_swiglu") - def original_bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False): - """Stub of Megatron's native impl. Returns Megatron's reference result so - fallback tests can assert numerics, and records that it ran.""" + def _swiglu(y): import torch import torch.nn.functional as F - original_bias_swiglu_impl.call_count += 1 - y = input if bias is None else input + bias y_1, y_2 = torch.chunk(y, 2, -1) return F.silu(y_1) * y_2 - original_bias_swiglu_impl.call_count = 0 + class BiasSwiGLUFunction: + """Stub of Megatron's bias autograd Function. Records that it ran.""" + + call_count = 0 + + @staticmethod + def apply(input, bias, fp8_input_store=False, cpu_offload_input=False): + BiasSwiGLUFunction.call_count += 1 + return _swiglu(input + bias) + + class SwiGLUFunction: + """Stub of Megatron's bias-free autograd Function -- the patch target.""" + + call_count = 0 + + @staticmethod + def apply(input, fp8_input_store=False, cpu_offload_input=False): + SwiGLUFunction.call_count += 1 + return _swiglu(input) + + class WeightedSwiGLUFunction: + """MoE routed-expert variant. Present so a test can assert Liger leaves it alone.""" + + @staticmethod + def apply(input, weights, fp8_input_store=False): + return _swiglu(input) * weights + def bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_input=False): + """Byte-for-byte reproduction of Megatron's dispatcher, including the global + lookups of ``BiasSwiGLUFunction`` / ``SwiGLUFunction`` that make the patch work.""" + ori_shape = input.shape + assert len(ori_shape) in [2, 3] + input = input.view(-1, ori_shape[-1]) + if bias is not None: + output = fused_swiglu.BiasSwiGLUFunction.apply(input, bias, fp8_input_store, cpu_offload_input) + else: + output = fused_swiglu.SwiGLUFunction.apply(input, fp8_input_store, cpu_offload_input) + return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) + + fused_swiglu.BiasSwiGLUFunction = BiasSwiGLUFunction + fused_swiglu.WeightedSwiGLUFunction = WeightedSwiGLUFunction + fused_swiglu.bias_swiglu_impl = bias_swiglu_impl if with_symbol: - fused_swiglu.bias_swiglu_impl = original_bias_swiglu_impl - # The MoE-weighted variant is deliberately out of scope for the current patch; it is - # present on the stub so a test can assert Liger leaves it alone. - fused_swiglu.weighted_bias_swiglu_impl = lambda *a, **kw: None + fused_swiglu.SwiGLUFunction = SwiGLUFunction sys.modules["megatron.core.fusions"] = fusions sys.modules["megatron.core.fusions.fused_bias_swiglu"] = fused_swiglu @@ -331,10 +320,9 @@ def original_bias_swiglu_impl(input, bias, fp8_input_store=False, cpu_offload_in shared_experts = types.ModuleType("megatron.core.transformer.moe.shared_experts") # Reproduce the by-name import: each consumer gets its own binding to the - # original function object. - if with_symbol: - mlp.bias_swiglu_impl = original_bias_swiglu_impl - shared_experts.bias_swiglu_impl = original_bias_swiglu_impl + # original ``bias_swiglu_impl`` object, captured before any patching. + mlp.bias_swiglu_impl = bias_swiglu_impl + shared_experts.bias_swiglu_impl = bias_swiglu_impl sys.modules["megatron.core.transformer"] = transformer sys.modules["megatron.core.transformer.moe"] = moe @@ -1169,10 +1157,14 @@ def test_rms_norm_only_patch_does_not_touch_ce_symbols(fake_megatron_ce): # =========================================================================== # 6. SwiGLU patch tests # =========================================================================== -# Liger replaces ``megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl`` — the symbol -# ``MLP.forward`` and ``SharedExpertMLP.forward`` call when bias_activation_fusion is on. -# Unlike CE and RMSNorm, its consumers use by-name imports, so the patch has to rewrite -# their module attributes too; several tests below exist specifically to enforce that. +# Liger replaces ``megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction`` -- the class +# Megatron's ``bias_swiglu_impl`` dispatches to when ``bias is None``, which is the path +# ``MLP.forward`` and ``SharedExpertMLP.forward`` take with bias_activation_fusion on. +# +# ``bias_swiglu_impl`` itself is deliberately *not* patched. It is a plain function that +# looks ``SwiGLUFunction`` up in its own module globals on every call, so replacing the +# class reaches consumers that already did ``from ... import bias_swiglu_impl`` -- no +# rebinding of consumer modules needed. Section 6.2 pins that property down. from liger_kernel.utils import infer_device # noqa: E402 @@ -1184,25 +1176,6 @@ def test_rms_norm_only_patch_does_not_touch_ce_symbols(fake_megatron_ce): ) -def _real_megatron_core_installed() -> bool: - """True when a genuine megatron-core is importable, as opposed to this file's stub. - - Evaluated once at import time, before any test installs a stub into ``sys.modules``. - The "missing megatron-core" tests can only assert anything meaningful when the package - really is absent -- which is the assumption in this repo's CI, but not in an - environment where someone installed megatron-core to run the e2e suite. - """ - import importlib.util - - try: - return importlib.util.find_spec("megatron.core.fusions.fused_bias_swiglu") is not None - except (ImportError, ValueError): - return False - - -_REAL_MEGATRON = _real_megatron_core_installed() - - # --------------------------------------------------------------------------- # 6.1 Symbol replacement + idempotency + opt-in. # --------------------------------------------------------------------------- @@ -1212,124 +1185,170 @@ def test_swiglu_patch_replaces_symbol(fake_megatron_swiglu): fused_swiglu, _ = fake_megatron_swiglu from liger_kernel.megatron import apply_liger_kernel_to_megatron - original = fused_swiglu.bias_swiglu_impl + original = fused_swiglu.SwiGLUFunction apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) - assert fused_swiglu.bias_swiglu_impl is not original - assert fused_swiglu.bias_swiglu_impl.__name__ == "liger_bias_swiglu_impl" - assert fused_swiglu.bias_swiglu_impl.__wrapped__ is original + assert fused_swiglu.SwiGLUFunction is not original + assert fused_swiglu.SwiGLUFunction.__wrapped__ is original def test_swiglu_patch_is_opt_in(fake_megatron_swiglu): - """``swiglu`` defaults to False — the symbol must be untouched unless asked for.""" - fused_swiglu, consumers = fake_megatron_swiglu + """``swiglu`` defaults to False -- the symbol must be untouched unless asked for.""" + fused_swiglu, _ = fake_megatron_swiglu from liger_kernel.megatron import apply_liger_kernel_to_megatron - before = fused_swiglu.bias_swiglu_impl + before = fused_swiglu.SwiGLUFunction apply_liger_kernel_to_megatron(rms_norm=False) - assert fused_swiglu.bias_swiglu_impl is before - assert consumers["mlp"].bias_swiglu_impl is before + assert fused_swiglu.SwiGLUFunction is before def test_swiglu_patch_is_idempotent(fake_megatron_swiglu): - """Calling apply twice must not stack wrappers — the sentinel attribute guards it.""" - fused_swiglu, consumers = fake_megatron_swiglu + """Calling apply twice must not stack wrappers -- the sentinel attribute guards it.""" + fused_swiglu, _ = fake_megatron_swiglu from liger_kernel.megatron import apply_liger_kernel_to_megatron apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) - first = fused_swiglu.bias_swiglu_impl - first_mlp = consumers["mlp"].bias_swiglu_impl + first = fused_swiglu.SwiGLUFunction apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) - assert fused_swiglu.bias_swiglu_impl is first - assert consumers["mlp"].bias_swiglu_impl is first_mlp + assert fused_swiglu.SwiGLUFunction is first # And the fallback chain must still be one deep, not wrapped around itself. - assert not hasattr(fused_swiglu.bias_swiglu_impl.__wrapped__, "__wrapped__") + assert not hasattr(fused_swiglu.SwiGLUFunction.__wrapped__, "__wrapped__") + +def test_swiglu_patch_leaves_sibling_functions_alone(fake_megatron_swiglu): + """Only the bias-free class is replaced. -def test_swiglu_patch_leaves_weighted_variant_alone(fake_megatron_swiglu): - """The MoE token-weighted variant is explicitly out of scope for this patch; it needs - a routing-weights gradient Liger has no kernel for.""" + ``BiasSwiGLUFunction`` has a bias term Liger's kernel does not implement, and + ``WeightedSwiGLUFunction`` (MoE routed experts) needs a routing-weights gradient Liger + has no kernel for. Leaving both untouched is what makes the bias/MoE fallbacks + structural rather than a runtime check. + """ fused_swiglu, _ = fake_megatron_swiglu from liger_kernel.megatron import apply_liger_kernel_to_megatron - before = fused_swiglu.weighted_bias_swiglu_impl + bias_before = fused_swiglu.BiasSwiGLUFunction + weighted_before = fused_swiglu.WeightedSwiGLUFunction + impl_before = fused_swiglu.bias_swiglu_impl + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) - assert fused_swiglu.weighted_bias_swiglu_impl is before + assert fused_swiglu.BiasSwiGLUFunction is bias_before + assert fused_swiglu.WeightedSwiGLUFunction is weighted_before + assert fused_swiglu.bias_swiglu_impl is impl_before # --------------------------------------------------------------------------- -# 6.2 Consumer-module rebinding — the failure mode unique to this patch. +# 6.2 Consumers with by-name imports are reached without being rebound. # --------------------------------------------------------------------------- -def test_swiglu_patch_rebinds_already_imported_consumers(fake_megatron_swiglu): +def test_swiglu_patch_reaches_consumers_without_rebinding_them(fake_megatron_swiglu): """``mlp.py`` and ``shared_experts.py`` do ``from ... import bias_swiglu_impl``. - If the patch only rewrites the defining module, both keep calling native Megatron and - the patch silently does nothing. This is the single most important SwiGLU-specific - assertion in this file. + Those bindings are stale copies the patch never touches -- and must not need to, + because ``bias_swiglu_impl`` resolves ``SwiGLUFunction`` from its defining module's + globals at call time. This is the assertion that justifies patching the class instead + of the function. """ fused_swiglu, consumers = fake_megatron_swiglu from liger_kernel.megatron import apply_liger_kernel_to_megatron + impl_before = fused_swiglu.bias_swiglu_impl apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) - patched = fused_swiglu.bias_swiglu_impl - assert consumers["mlp"].bias_swiglu_impl is patched - assert consumers["shared_experts"].bias_swiglu_impl is patched + # Consumers still hold Megatron's original function object... + assert consumers["mlp"].bias_swiglu_impl is impl_before + assert consumers["shared_experts"].bias_swiglu_impl is impl_before + assert fused_swiglu.SwiGLUFunction is not fused_swiglu.SwiGLUFunction.__wrapped__ + # ...yet a call through those stale bindings still resolves whatever + # ``fused_bias_swiglu.SwiGLUFunction`` currently is. A probe stands in for Liger here + # so the assertion holds on CPU, where the Triton kernel cannot run. + import torch + + class _Probe: + called = 0 + + @staticmethod + def apply(input, fp8_input_store=False, cpu_offload_input=False): + _Probe.called += 1 + return input[..., : input.shape[-1] // 2] + + fused_swiglu.SwiGLUFunction = _Probe + for consumer in consumers.values(): + consumer.bias_swiglu_impl(torch.zeros(8, 128), None, False, False) + + assert _Probe.called == len(consumers) -def test_swiglu_patch_tolerates_absent_consumer_modules(fake_megatron_swiglu_no_consumers): - """Consumers not yet imported are fine — they'll import the already-patched symbol. - The patch must not crash trying to rebind them.""" + +def test_swiglu_patch_is_import_order_independent(fake_megatron_swiglu_no_consumers): + """Patching before consumers exist works too, for the same reason. + + The old failure mode -- patch applied, log emitted, nothing actually swapped -- was + entirely a function of whether the consumer had imported the symbol yet. Under the + class-level patch neither order matters, and the patch has nothing to rebind so it + cannot crash on absent consumer modules. + """ fused_swiglu = fake_megatron_swiglu_no_consumers from liger_kernel.megatron import apply_liger_kernel_to_megatron + original = fused_swiglu.SwiGLUFunction apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + patched = fused_swiglu.SwiGLUFunction + assert patched is not original - assert fused_swiglu.bias_swiglu_impl.__name__ == "liger_bias_swiglu_impl" + # A consumer importing *after* the patch picks up the dispatcher, which resolves the + # class lazily -- so it lands on the patched object, not on the pre-patch one. + import torch + late_consumer = types.ModuleType("late_consumer") + late_consumer.bias_swiglu_impl = fused_swiglu.bias_swiglu_impl -def test_swiglu_patch_does_not_clobber_unrelated_consumer_binding(fake_megatron_swiglu): - """If a consumer's binding isn't the original (a user already replaced it), leave it.""" - fused_swiglu, consumers = fake_megatron_swiglu - from liger_kernel.megatron import apply_liger_kernel_to_megatron + class _Probe: + called = 0 - sentinel = object() - consumers["mlp"].bias_swiglu_impl = sentinel + @staticmethod + def apply(input, fp8_input_store=False, cpu_offload_input=False): + _Probe.called += 1 + return input[..., : input.shape[-1] // 2] - apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + fused_swiglu.SwiGLUFunction = _Probe + late_consumer.bias_swiglu_impl(torch.zeros(8, 128), None, False, False) - assert consumers["mlp"].bias_swiglu_impl is sentinel - # The other consumer, still on the original, is rebound normally. - assert consumers["shared_experts"].bias_swiglu_impl is fused_swiglu.bias_swiglu_impl + assert _Probe.called == 1 # --------------------------------------------------------------------------- -# 6.3 Missing megatron-core / missing symbol → actionable ImportError. +# 6.3 Missing megatron-core / missing symbol -> actionable ImportError. # --------------------------------------------------------------------------- -@pytest.mark.skipif(_REAL_MEGATRON, reason="a real megatron-core is installed, so the import cannot fail") -def test_swiglu_patch_raises_when_megatron_missing(): - from liger_kernel.megatron import apply_liger_kernel_to_megatron - +def test_swiglu_patch_raises_when_megatron_not_installed(): _uninstall_fake_megatron() - with pytest.raises(ImportError, match="fused_bias_swiglu"): - apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + real_import = __builtins__["__import__"] if isinstance(__builtins__, dict) else __builtins__.__import__ + + def blocking_import(name, *args, **kwargs): + if name == "megatron" or name.startswith("megatron."): + raise ImportError(f"No module named {name!r}") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=blocking_import): + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + with pytest.raises(ImportError, match="requires megatron-core"): + apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) def test_swiglu_patch_raises_when_symbol_missing(): """A renamed symbol in a future megatron-core must fail loudly, not silently no-op.""" - from liger_kernel.megatron import apply_liger_kernel_to_megatron - _install_fake_megatron_swiglu(with_symbol=False) try: - with pytest.raises(ImportError, match="bias_swiglu_impl not found"): + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + with pytest.raises(ImportError, match="symbol path may have changed"): apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) finally: _uninstall_fake_megatron() @@ -1349,17 +1368,20 @@ def _swiglu_reference(y): @_requires_accelerator -def test_patched_swiglu_symbol_computes_correct_output(fake_megatron_swiglu): +@pytest.mark.parametrize("shape", [(8, 128), (2, 4, 128)]) +def test_patched_swiglu_symbol_computes_correct_output(fake_megatron_swiglu, shape): + """3D input is covered because Megatron's dispatcher owns the flatten/restore, and + the patch relies on that rather than reimplementing it.""" import torch from liger_kernel.megatron import apply_liger_kernel_to_megatron from liger_kernel.utils import infer_device - fused_swiglu, consumers = fake_megatron_swiglu + _, consumers = fake_megatron_swiglu apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) - x = torch.randn(8, 128, device=infer_device(), dtype=torch.float32) - # Call through the *consumer's* binding, exactly as MLP.forward does — positionally. + x = torch.randn(*shape, device=infer_device(), dtype=torch.float32) + # Call through the *consumer's* binding, exactly as MLP.forward does -- positionally. out = consumers["mlp"].bias_swiglu_impl(x, None, False, False) torch.testing.assert_close(out, _swiglu_reference(x), atol=1e-5, rtol=1e-5) @@ -1386,23 +1408,26 @@ def test_patched_swiglu_symbol_preserves_gradients(fake_megatron_swiglu): torch.testing.assert_close(h_liger.grad, h_ref.grad, atol=1e-5, rtol=1e-5) -def test_patched_swiglu_symbol_falls_back_for_bias(fake_megatron_swiglu): - """Bias is added pre-activation in Megatron and Liger has no bias term, so the patch - must hand the call straight back to the original. Runs on CPU — no Triton involved, - which is the point: the fallback path never touches the kernel.""" +def test_patched_swiglu_leaves_bias_path_on_megatron(fake_megatron_swiglu): + """Bias is added pre-activation in Megatron and Liger has no bias term. + + With the class-level patch this needs no runtime check: a non-``None`` bias is routed + to ``BiasSwiGLUFunction`` by Megatron's own dispatcher, which we never replace. Runs on + CPU precisely because the kernel is never reached. + """ import torch from liger_kernel.megatron import apply_liger_kernel_to_megatron fused_swiglu, consumers = fake_megatron_swiglu - original = fused_swiglu.bias_swiglu_impl apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) x = torch.randn(8, 128) bias = torch.randn(128) out = consumers["mlp"].bias_swiglu_impl(x, bias, False, False) - assert original.call_count == 1, "expected the native impl to be invoked for the bias path" + assert fused_swiglu.BiasSwiGLUFunction.call_count == 1 + assert fused_swiglu.SwiGLUFunction.__wrapped__.call_count == 0 torch.testing.assert_close(out, _swiglu_reference(x + bias), atol=1e-5, rtol=1e-5) @@ -1413,13 +1438,15 @@ def test_patched_swiglu_symbol_falls_back_for_bias(fake_megatron_swiglu): def test_patched_swiglu_symbol_falls_back_for_unsupported_flags( fake_megatron_swiglu, fp8_input_store, cpu_offload_input ): + """fp8 input storage and CPU offload are activation-storage policies Liger's kernel + does not implement, so those calls go back to the wrapped original.""" import torch from liger_kernel.megatron import apply_liger_kernel_to_megatron fused_swiglu, consumers = fake_megatron_swiglu - original = fused_swiglu.bias_swiglu_impl apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) + original = fused_swiglu.SwiGLUFunction.__wrapped__ x = torch.randn(8, 128) consumers["mlp"].bias_swiglu_impl(x, None, fp8_input_store, cpu_offload_input) @@ -1441,11 +1468,12 @@ def test_swiglu_patch_does_not_touch_other_kernels(fake_megatron_swiglu): ce_before = fused_ce.fused_vocab_parallel_cross_entropy unfused_before = unfused_ce.vocab_parallel_cross_entropy + swiglu_before = fused_swiglu.SwiGLUFunction apply_liger_kernel_to_megatron(rms_norm=False, swiglu=True) assert fused_ce.fused_vocab_parallel_cross_entropy is ce_before assert unfused_ce.vocab_parallel_cross_entropy is unfused_before - assert fused_swiglu.bias_swiglu_impl.__name__ == "liger_bias_swiglu_impl" + assert fused_swiglu.SwiGLUFunction is not swiglu_before finally: _uninstall_fake_megatron() diff --git a/test/megatron/test_swiglu.py b/test/megatron/test_swiglu.py index 8a500abba..26893ec87 100644 --- a/test/megatron/test_swiglu.py +++ b/test/megatron/test_swiglu.py @@ -1,22 +1,12 @@ """Unit tests for LigerMegatronSwiGLU. -Sections 1-5 deliberately do not import megatron-core; the wrapper's contract is verified +These tests deliberately do not import megatron-core; the wrapper's contract is verified against a local reproduction of Megatron's own reference:: y_1, y_2 = torch.chunk(y, 2, -1) return F.silu(y_1) * y_2 The parametrization style mirrors ``test/megatron/test_rms_norm.py``. - -Section 6 is different, and deliberately so. Verifying only against a reimplementation -leaves one class of bug completely uncovered: anything where our *belief* about Megatron's -API is wrong. A signature change, a renamed module path, or a different dispatch condition -would leave every test above green while the integration silently breaks at runtime. Those -tests import the real megatron-core and skip cleanly when it is absent, so they cost -nothing in the default CI environment. Run them with:: - - pip install megatron-core - pytest test/megatron/test_swiglu.py -v """ import os @@ -390,253 +380,3 @@ def test_signature_matches_megatron_positional_order(): params = list(inspect.signature(LigerMegatronSwiGLU.forward).parameters) assert params == ["self", "input", "bias", "fp8_input_store", "cpu_offload_input"] - - -# --------------------------------------------------------------------------- -# 6. End-to-end validation against a real megatron-core install -# --------------------------------------------------------------------------- -# -# Everything above verifies the wrapper against our own reproduction of Megatron's -# formula. That cannot catch a wrong belief about Megatron's API. The tests below import -# the real package and pin: the true signature and positional order of -# ``bias_swiglu_impl``, numerical parity with Megatron's own TorchScript implementation, -# that the consumer module paths the monkey patch rebinds actually exist and are rebound, -# and that a real ``MLP`` produces identical output and gradients when patched. -# -# The import is conditional rather than a module-level ``pytest.importorskip`` so that the -# dependency-free tests above still run when megatron-core is absent. - -try: - import megatron.core # noqa: F401 - - _MEGATRON_AVAILABLE = True -except ImportError: - _MEGATRON_AVAILABLE = False - -requires_megatron = pytest.mark.skipif( - not _MEGATRON_AVAILABLE - or not (torch.cuda.is_available() or (getattr(torch, "xpu", None) and torch.xpu.is_available())), - reason="requires megatron-core and an accelerator", -) - -_PATCH_TARGETS = ( - "megatron.core.fusions.fused_bias_swiglu", - "megatron.core.transformer.mlp", - "megatron.core.transformer.moe.shared_experts", -) - - -@pytest.fixture -def restore_patch(): - """Snapshot and restore ``bias_swiglu_impl`` in every module the patch touches. - - Restoring only the consumer modules is not enough: the patch is idempotent on the - *defining* module, so leaving that one patched makes a later ``apply_...`` call - short-circuit and the next test observes a half-applied patch. - """ - import importlib - - mods = [importlib.import_module(name) for name in _PATCH_TARGETS] - saved = [(m, m.bias_swiglu_impl) for m in mods] - try: - yield {m.__name__: fn for m, fn in saved} - finally: - for m, fn in saved: - m.bias_swiglu_impl = fn - - -@requires_megatron -def test_real_bias_swiglu_impl_signature_is_what_the_wrapper_assumes(): - """``LigerMegatronSwiGLU.forward`` mirrors this signature positionally. - - Both ``MLP.forward`` and ``SharedExpertMLP.forward`` call ``bias_swiglu_impl`` - positionally, so parameter *order* is part of the contract, not just parameter names. - If Megatron reorders or inserts an argument this fails loudly instead of the patch - silently passing ``fp8_input_store`` where ``bias`` was expected. - """ - import inspect - - from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl - - real = list(inspect.signature(bias_swiglu_impl).parameters) - ours = [p for p in inspect.signature(LigerMegatronSwiGLU.forward).parameters if p != "self"] - - assert real[: len(ours)] == ours, ( - f"megatron-core's bias_swiglu_impl signature {real} no longer matches the order " - f"LigerMegatronSwiGLU.forward accepts {ours}." - ) - - -@requires_megatron -@pytest.mark.parametrize("dtype, atol, rtol", [(torch.float32, 1e-6, 1e-6), (torch.bfloat16, 1e-2, 1e-2)]) -@pytest.mark.parametrize("shape", [(4, 2, 512), (2, 1, 11008)]) -def test_matches_real_megatron_numerically(shape, dtype, atol, rtol): - """Forward and backward parity against Megatron's actual TorchScript implementation.""" - from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl - - s, b, ffn = shape - _tensor = torch.randn(s, b, 2 * ffn, device=device, dtype=dtype) - do = torch.randn(s, b, ffn, device=device, dtype=dtype) - - h_ref = _tensor.clone().requires_grad_(True) - ref = bias_swiglu_impl(h_ref, None, False, False) - ref.backward(do) - - h_liger = _tensor.clone().requires_grad_(True) - got = LigerMegatronSwiGLU()(h_liger, None, False, False) - got.backward(do) - - assert got.shape == ref.shape - torch.testing.assert_close(got, ref, atol=atol, rtol=rtol) - torch.testing.assert_close(h_liger.grad, h_ref.grad, atol=atol, rtol=rtol) - - -@requires_megatron -def test_patch_rebinds_the_real_consumer_modules(restore_patch): - """The by-name imports in the real megatron-core must actually be rebound. - - ``megatron.core.transformer.mlp`` does ``from ...fused_bias_swiglu import - bias_swiglu_impl`` at import time, so patching only the defining module leaves the - consumer holding the original function object. This verifies both that those module - paths still exist and that the rebinding reaches them. - """ - import importlib - - from liger_kernel.megatron.monkey_patch import apply_liger_kernel_to_megatron - - apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=False, swiglu=True) - - for name in _PATCH_TARGETS: - mod = importlib.import_module(name) - patched = mod.bias_swiglu_impl - original = restore_patch[name] - assert patched is not original, f"{name}.bias_swiglu_impl was not rebound" - # The patch installs a plain function wrapping a shared LigerMegatronSwiGLU and - # tags it so re-application is idempotent. __wrapped__ must point back at - # Megatron's real function so unsupported configs fall back rather than recurse. - assert getattr(patched, "__liger_patched__", False), f"{name} rebound to an untagged object" - assert patched.__wrapped__ is original - - # And it must actually compute Liger's result, not just look patched. - mlp_mod = importlib.import_module("megatron.core.transformer.mlp") - x = torch.randn(4, 2, 128, device=device, dtype=torch.float32) - torch.testing.assert_close( - mlp_mod.bias_swiglu_impl(x, None, False, False), - LigerMegatronSwiGLU()(x, None, False, False), - atol=0, - rtol=0, - ) - - -@requires_megatron -def test_patch_rebinds_consumers_imported_after_the_first_call(restore_patch): - """Re-applying the patch must still reach a consumer that reverted to the original. - - Real-world shape of this: a consumer module is imported lazily, *after* - ``apply_liger_kernel_to_megatron`` already ran, so it captures Megatron's original at - its own import time. A naive "already patched, return" short-circuit would leave that - consumer permanently unpatched. - """ - import importlib - - from liger_kernel.megatron.monkey_patch import apply_liger_kernel_to_megatron - - apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=False, swiglu=True) - - mlp_mod = importlib.import_module("megatron.core.transformer.mlp") - mlp_mod.bias_swiglu_impl = restore_patch["megatron.core.transformer.mlp"] - - apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=False, swiglu=True) - - assert getattr(mlp_mod.bias_swiglu_impl, "__liger_patched__", False), ( - "a consumer that reverted to Megatron's original was not re-bound on re-apply" - ) - - -@pytest.fixture(scope="module") -def model_parallel(): - """Minimal single-rank model-parallel init so real Megatron layers can be built.""" - from megatron.core import parallel_state - from megatron.core import tensor_parallel - - os.environ.setdefault("MASTER_ADDR", "127.0.0.1") - os.environ.setdefault("MASTER_PORT", "29591") - os.environ.setdefault("RANK", "0") - os.environ.setdefault("WORLD_SIZE", "1") - - if not torch.distributed.is_initialized(): - torch.distributed.init_process_group(backend="nccl", world_size=1, rank=0) - parallel_state.initialize_model_parallel(tensor_model_parallel_size=1) - # ColumnParallelLinear initializes its weight under the model-parallel RNG fork, which - # does not exist until this is called. Without it, building any parallel layer raises - # "cuda rng state model-parallel-rng is not added". - tensor_parallel.model_parallel_cuda_manual_seed(0) - torch.manual_seed(0) - yield - parallel_state.destroy_model_parallel() - - -def _build_mlp(hidden, ffn, dtype): - from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec - from megatron.core.transformer.mlp import MLP - from megatron.core.transformer.transformer_config import TransformerConfig - - config = TransformerConfig( - num_layers=1, - hidden_size=hidden, - ffn_hidden_size=ffn, - num_attention_heads=4, - gated_linear_unit=True, - activation_func=F.silu, - bias_activation_fusion=True, - add_bias_linear=False, - bf16=(dtype is torch.bfloat16), - params_dtype=dtype, - ) - spec = get_gpt_layer_local_spec().submodules.mlp.submodules - return MLP(config, spec).to(device=device, dtype=dtype) - - -@requires_megatron -def test_real_mlp_forward_backward_parity(model_parallel, restore_patch): - """A real Megatron ``MLP`` must produce identical results with the patch applied. - - This is the test that exercises the actual dispatch path -- ``MLP.forward`` deciding to - call ``bias_swiglu_impl`` based on ``config.bias_activation_fusion`` / - ``gated_linear_unit`` / ``activation_func``, with the real ``linear_fc1`` output layout - feeding it. - - One MLP instance is reused for both runs rather than building two and copying weights: - the patch is a module-level symbol swap, so the same instance picks it up, and this - sidesteps Megatron's ``_extra_state`` entries (which are None and break a naive - ``state_dict`` round-trip). - """ - import importlib - - from liger_kernel.megatron.monkey_patch import apply_liger_kernel_to_megatron - - hidden, ffn, dtype = 128, 512, torch.float32 - mlp_mod = importlib.import_module("megatron.core.transformer.mlp") - original = restore_patch["megatron.core.transformer.mlp"] - - mlp = _build_mlp(hidden, ffn, dtype) - x = torch.randn(8, 2, hidden, device=device, dtype=dtype) - - def run(): - for p_ in mlp.parameters(): - p_.grad = None - h = x.clone().requires_grad_(True) - out, _ = mlp(h) - out.sum().backward() - return out.detach().clone(), h.grad.clone(), [p_.grad.clone() for p_ in mlp.parameters()] - - ref_out, ref_dx, ref_dw = run() - - apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=False, swiglu=True) - assert mlp_mod.bias_swiglu_impl is not original, "patch did not take effect" - liger_out, liger_dx, liger_dw = run() - - torch.testing.assert_close(liger_out, ref_out, atol=1e-5, rtol=1e-5) - torch.testing.assert_close(liger_dx, ref_dx, atol=1e-5, rtol=1e-5) - for got, want in zip(liger_dw, ref_dw): - torch.testing.assert_close(got, want, atol=1e-5, rtol=1e-5)