@@ -320,6 +320,7 @@ def _(
320320 "ademamix" : ADEMAMIX ,
321321}
322322
323+
323324@torch .compile
324325def _optimizer_precondition_32bit (
325326 g : torch .Tensor ,
@@ -525,29 +526,53 @@ def _(
525526
526527 if optimizer_name == "lion" :
527528 _optimizer_update_32bit (
528- g , p , state1 , state2 , unorm_vec , max_unorm , param_norm ,
529- beta1 , beta2 , beta3 , alpha , eps , weight_decay , step ,
530- lr , gnorm_scale , optimizer_id
529+ g ,
530+ p ,
531+ state1 ,
532+ state2 ,
533+ unorm_vec ,
534+ max_unorm ,
535+ param_norm ,
536+ beta1 ,
537+ beta2 ,
538+ beta3 ,
539+ alpha ,
540+ eps ,
541+ weight_decay ,
542+ step ,
543+ lr ,
544+ gnorm_scale ,
545+ optimizer_id ,
531546 )
532547
533548 if max_unorm > 0.0 :
534549 unorm_vec .zero_ ()
535550 _optimizer_precondition_32bit (
536- g , p , state1 , state2 , unorm_vec ,
537- beta1 , beta2 , eps , weight_decay , step ,
538- lr , gnorm_scale , optimizer_id
551+ g , p , state1 , state2 , unorm_vec , beta1 , beta2 , eps , weight_decay , step , lr , gnorm_scale , optimizer_id
539552 )
540553 else :
541554 if max_unorm > 0.0 :
542555 unorm_vec .zero_ ()
543556 _optimizer_precondition_32bit (
544- g , p , state1 , state2 , unorm_vec ,
545- beta1 , beta2 , eps , weight_decay , step ,
546- lr , gnorm_scale , optimizer_id
557+ g , p , state1 , state2 , unorm_vec , beta1 , beta2 , eps , weight_decay , step , lr , gnorm_scale , optimizer_id
547558 )
548559
549560 _optimizer_update_32bit (
550- g , p , state1 , state2 , unorm_vec , max_unorm , param_norm ,
551- beta1 , beta2 , beta3 , alpha , eps , weight_decay , step ,
552- lr , gnorm_scale , optimizer_id
561+ g ,
562+ p ,
563+ state1 ,
564+ state2 ,
565+ unorm_vec ,
566+ max_unorm ,
567+ param_norm ,
568+ beta1 ,
569+ beta2 ,
570+ beta3 ,
571+ alpha ,
572+ eps ,
573+ weight_decay ,
574+ step ,
575+ lr ,
576+ gnorm_scale ,
577+ optimizer_id ,
553578 )
0 commit comments