perforatedai.utils_perforatedai
1# Copyright (c) 2025 Perforated AI 2 3import torch 4import torch.nn as nn 5import torch.nn.init as init 6import torch.nn.functional as F 7import math 8import sys 9import numpy as np 10import pdb 11import os 12import time 13import warnings 14from collections import defaultdict 15 16from perforatedai import globals_perforatedai as GPA 17from perforatedai import modules_perforatedai as PA 18from perforatedai import tracker_perforatedai as TPA 19from perforatedai import clean_perforatedai as CL 20from perforatedai import blockwise_perforatedai as BPA 21from perforatedai import network_perforatedai as NPA 22 23try: 24 from dashboard_utils.event_emitter import emitter as _dashboard_emitter 25except ImportError: 26 _dashboard_emitter = None 27 28try: 29 from perforatedbp import utils_pbp as UPB 30 from perforatedbp import modules_pbp as MPB 31except ModuleNotFoundError as e: 32 # Only pass if perforatedbp package itself is missing 33 if e.name == "perforatedbp": 34 pass 35 else: 36 # perforatedbp exists but is missing a dependency 37 raise 38 39import copy 40 41from safetensors.torch import load_file 42from safetensors.torch import save_file 43from safetensors.torch import safe_open 44 45 46def perforate_model( 47 model, 48 doing_pai=True, 49 save_name="", 50 making_graphs=True, 51 maximizing_score=True, 52 num_classes=10000000000, 53 values_per_train_epoch=-1, 54 values_per_val_epoch=-1, 55 zooming_graph=True, 56): 57 """Main function to initialize the network to add dendrites 58 59 This kicks off the entire Perforated AI process to add 60 the scaffolding to the network to be able to add dendrites 61 62 Parameters 63 ---------- 64 model : nn.Module 65 The neural network model to initialize. 66 doing_pai : bool, optional 67 Whether to actually add dendrites, by default True 68 save_name : str, optional 69 The name to save the model under, by default "PAI" 70 making_graphs : bool, optional 71 Whether to create graphs during training, by default True 72 maximizing_score : bool, optional 73 Whether to maximize the score during training, by default True 74 setting to false is for when the score is a loss to be minimized 75 num_classes : int, optional 76 The number of output classes, unused in current version 77 values_per_train_epoch : int, optional 78 The number of values to look back for graphing 79 during training, by default -1 (all values). 80 values_per_val_epoch : int, optional 81 The number of values to look back for graphing 82 during validation, by default -1 (all values). 83 zooming_graph : bool, optional 84 Whether to enable zooming on the graphs, by default True 85 86 Returns 87 ------- 88 model : nn.Module 89 The modified model with dendrite scaffolding added if doing_pai is True 90 91 """ 92 93 if save_name == "": 94 if GPA.pc.get_save_name() == "": 95 save_name = "PAI" 96 else: 97 save_name = GPA.pc.get_save_name() 98 99 if "/" in save_name: 100 print( 101 f"Warning: save_name '{save_name}' contains '/'. Relative paths are not implemented yet." 102 ) 103 sys.exit(1) 104 105 sanitized_save_name = "".join( 106 ch for ch in save_name if ch.isalnum() or ch in ("_", "-", ".") 107 ) 108 if sanitized_save_name != save_name: 109 print( 110 f"Warning: save_name '{save_name}' contained spaces or special characters. " 111 f"Using '{sanitized_save_name}' instead." 112 ) 113 save_name = sanitized_save_name 114 115 if save_name == "": 116 print("Warning: save_name became empty after sanitization. Using 'PAI'.") 117 save_name = "PAI" 118 119 120 GPA.pai_tracker = TPA.PAINeuronModuleTracker( 121 doing_pai=doing_pai, save_name=save_name 122 ) 123 GPA.pc.set_save_name(save_name) 124 if _dashboard_emitter is not None: 125 _dashboard_emitter.emit_run_start(GPA.pc, save_name) 126 model = GPA.pai_tracker.initialize( 127 model, 128 doing_pai=doing_pai, 129 save_name=save_name, 130 making_graphs=making_graphs, 131 maximizing_score=maximizing_score, 132 num_classes=num_classes, 133 values_per_train_epoch=-values_per_train_epoch, 134 values_per_val_epoch=values_per_val_epoch, 135 zooming_graph=zooming_graph, 136 ) 137 138 # Save config after perforation 139 if not GPA.pc.get_testing_dendrite_capacity(): 140 import os 141 GPA.pc.save_config(os.path.join(os.getcwd(), save_name, f"{save_name}_config.json")) 142 143 return model 144 145 146def get_pai_modules(net, depth, seen_ids=None): 147 """Get a list of all neuron modules 148 149 Parameters 150 ---------- 151 net : nn.Module 152 The module to search. 153 depth : int 154 The current depth in the recursion. 155 156 Returns 157 ------- 158 list 159 A list of all PAI neuron modules found in the network. 160 161 """ 162 if seen_ids is None: 163 seen_ids = set() 164 all_members = net.__dir__() 165 this_list = [] 166 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 167 for submodule_id, layer in net.named_children(): 168 # If there is a self pointer ignore it 169 if net.get_submodule(submodule_id) is net: 170 continue 171 if type(net.get_submodule(submodule_id)) is PA.PAINeuronModule: 172 module = net.get_submodule(submodule_id) 173 if id(module) in seen_ids: 174 continue 175 seen_ids.add(id(module)) 176 this_list = this_list + [module] 177 else: 178 this_list = this_list + get_pai_modules( 179 net.get_submodule(submodule_id), depth + 1, seen_ids 180 ) 181 else: 182 for member in all_members: 183 if isinstance(getattr(type(net), member, None), property): 184 continue 185 # if the getter fails or it is a self pointer ignore it 186 try: 187 if getattr(net, member, None) is net: 188 continue 189 except: 190 continue 191 if type(getattr(net, member, None)) is PA.PAINeuronModule: 192 module = getattr(net, member) 193 if id(module) in seen_ids: 194 continue 195 seen_ids.add(id(module)) 196 this_list = this_list + [module] 197 elif ( 198 issubclass(type(getattr(net, member, None)), nn.Module) 199 or issubclass(type(getattr(net, member, None)), nn.Sequential) 200 or issubclass(type(getattr(net, member, None)), nn.ModuleList) 201 ): 202 this_list = this_list + get_pai_modules( 203 getattr(net, member), depth + 1, seen_ids 204 ) 205 206 return this_list 207 208 209def get_tracked_modules(net, depth, seen_ids=None): 210 """Get a list of all tracked modules 211 212 Parameters 213 ---------- 214 net : nn.Module 215 The module to search. 216 depth : int 217 The current depth in the recursion. 218 219 Returns 220 ------- 221 list 222 A list of all tracked modules found in the network. 223 224 """ 225 if seen_ids is None: 226 seen_ids = set() 227 all_members = net.__dir__() 228 this_list = [] 229 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 230 for submodule_id, layer in net.named_children(): 231 if net.get_submodule(submodule_id) is net: 232 continue 233 if type(net.get_submodule(submodule_id)) is PA.TrackedNeuronModule: 234 module = net.get_submodule(submodule_id) 235 if id(module) in seen_ids: 236 continue 237 seen_ids.add(id(module)) 238 this_list = this_list + [module] 239 else: 240 this_list = this_list + get_tracked_modules( 241 net.get_submodule(submodule_id), depth + 1, seen_ids 242 ) 243 else: 244 for member in all_members: 245 if isinstance(getattr(type(net), member, None), property): 246 continue 247 # if the getter fails or it is a self pointer ignore it 248 try: 249 if getattr(net, member, None) is net: 250 continue 251 except: 252 continue 253 if type(getattr(net, member, None)) is PA.TrackedNeuronModule: 254 module = getattr(net, member) 255 if id(module) in seen_ids: 256 continue 257 seen_ids.add(id(module)) 258 this_list = this_list + [module] 259 elif issubclass(type(getattr(net, member, None)), nn.Module): 260 this_list = this_list + get_tracked_modules( 261 getattr(net, member), depth + 1, seen_ids 262 ) 263 return this_list 264 265 266def get_pai_module_params(net, depth, seen_ids=None): 267 """Get a list of all neuron module parameters 268 269 Parameters 270 ---------- 271 net : nn.Module 272 The module to search. 273 depth : int 274 The current depth in the recursion. 275 276 Returns 277 ------- 278 list 279 A list of all parameters of neuron modules found in this module. 280 281 """ 282 283 if seen_ids is None: 284 seen_ids = set() 285 all_members = net.__dir__() 286 this_list = [] 287 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 288 for submodule_id, layer in net.named_children(): 289 if isinstance(net.get_submodule(submodule_id), PA.PAINeuronModule): # 290 module = net.get_submodule(submodule_id) 291 if id(module) in seen_ids: 292 continue 293 seen_ids.add(id(module)) 294 for param in module.parameters(): 295 if param.requires_grad: 296 this_list = this_list + [param] 297 else: 298 this_list = this_list + get_pai_module_params( 299 net.get_submodule(submodule_id), depth + 1, seen_ids 300 ) 301 else: 302 for member in all_members: 303 if isinstance(getattr(type(net), member, None), property): 304 continue 305 if getattr(net, member, None) == net: 306 continue 307 if isinstance(getattr(net, member, None), PA.PAINeuronModule): 308 module = getattr(net, member) 309 if id(module) in seen_ids: 310 continue 311 seen_ids.add(id(module)) 312 for param in module.parameters(): 313 if param.requires_grad: 314 this_list = this_list + [param] 315 elif issubclass(type(getattr(net, member, None)), nn.Module): 316 this_list = this_list + get_pai_module_params( 317 getattr(net, member), depth + 1, seen_ids 318 ) 319 return this_list 320 321 322def get_pai_network_params(net): 323 """Get a list of all neuron module parameters 324 325 Parameters 326 ---------- 327 net : nn.Module 328 The full model to search. 329 330 Returns 331 ------- 332 list 333 A list of all parameters of neuron modules found in the network. 334 335 """ 336 param_list = get_pai_module_params(net, 0) 337 return param_list 338 339 340def replace_predefined_modules(start_module): 341 """Replace a module with the module from globals list 342 343 Parameters 344 ---------- 345 start_module : nn.Module 346 The module to replace. 347 348 Returns 349 ------- 350 nn.Module 351 The replaced module. 352 353 """ 354 index = GPA.pc.get_modules_to_replace().index(type(start_module)) 355 return GPA.pc.get_replacement_modules()[index](start_module) 356 357 358def scan_module_aliases(net): 359 """Find alias module paths that point to already-seen module instances. 360 361 Parameters 362 ---------- 363 net : Any PyTorch Module. 364 365 Returns 366 ------- 367 dict[str, str] 368 Mapping of alias module paths to their canonical module paths. 369 """ 370 canonical = {} 371 aliases = {} 372 for name, module in net.named_modules(remove_duplicate=False): 373 if name == "": 374 continue 375 sub_name = "." + name 376 module_id = id(module) 377 if module_id in canonical: 378 aliases[sub_name] = canonical[module_id] 379 else: 380 canonical[module_id] = sub_name 381 return aliases 382 383 384def convert_module( 385 net, 386 depth, 387 name_so_far, 388 converted_list, 389 converted_names_list, 390 neuron_module_class, 391 tracked_module_class, 392): 393 """Recursive function to do all conversion of modules to wrappers of modules 394 395 This is the function that goes through all of the module lists from 396 the globals file and does all the conversion and replacements to 397 setup the dendrite scaffolding as instructed. 398 399 Parameters 400 ---------- 401 net : nn.Module 402 The module to convert. 403 depth : int 404 The current depth in the recursion. 405 name_so_far : str 406 The name of the module so far in the recursion. 407 converted_list : list 408 A list of already converted module ids to avoid infinite loops. 409 converted_names_list : list 410 A corresponding list to help debug duplicate conversions 411 412 Returns 413 ------- 414 nn.Module 415 The converted module. 416 417 """ 418 if GPA.pc.get_verbose(): 419 print("calling convert on %s depth %d" % (net, depth)) 420 print( 421 "calling convert on %s: %s, depth %d" 422 % (name_so_far, type(net).__name__, depth) 423 ) 424 if isinstance(net, neuron_module_class) or ( 425 (tracked_module_class is not None) and isinstance(net, tracked_module_class) 426 ): 427 if GPA.pc.get_verbose(): 428 print( 429 "This is only being called because something in your model " 430 "is pointed to twice by two different variables. Highest " 431 "thing on the list is one of the duplicates" 432 ) 433 return net 434 if depth == 0 and name_so_far == "": 435 aliases = scan_module_aliases(net) 436 existing_not_save = set(GPA.pc.get_module_names_to_not_save()) 437 aliases_to_skip = [ 438 alias for alias in aliases.keys() if alias not in existing_not_save 439 ] 440 if aliases_to_skip: 441 GPA.pc.append_module_names_to_not_save(aliases_to_skip) 442 print( 443 "Auto-detected duplicate module aliases via named_modules; " 444 "keeping first-seen paths and skipping:" 445 ) 446 for alias in aliases_to_skip: 447 print(" - %s (keeps %s)" % (alias, aliases[alias])) 448 all_members = net.__dir__() 449 if GPA.pc.get_extra_verbose(): 450 print("all members:") 451 for member in all_members: 452 print(" - %s" % member) 453 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 454 for submodule_id, layer in net.named_children(): 455 sub_name = name_so_far + "." + str(submodule_id) 456 if sub_name in GPA.pc.get_module_ids_to_track(): 457 if GPA.pc.get_verbose(): 458 print("Seq ID is in track IDs: %s" % sub_name) 459 if tracked_module_class is None: 460 continue 461 setattr( 462 net, 463 submodule_id, 464 tracked_module_class(net.get_submodule(submodule_id), sub_name), 465 ) 466 continue 467 if sub_name in GPA.pc.get_module_ids_to_perforate(): 468 if GPA.pc.get_verbose(): 469 print("Seq ID is in convert IDs: %s" % sub_name) 470 setattr( 471 net, 472 submodule_id, 473 neuron_module_class(net.get_submodule(submodule_id), sub_name), 474 ) 475 continue 476 if type(net.get_submodule(submodule_id)) in GPA.pc.get_modules_to_replace(): 477 if GPA.pc.get_verbose(): 478 print( 479 "Seq sub is in replacement module so replacing: %s" % sub_name 480 ) 481 setattr( 482 net, 483 submodule_id, 484 replace_predefined_modules(net.get_submodule(submodule_id)), 485 ) 486 if ( 487 type(net.get_submodule(submodule_id)) in GPA.pc.get_modules_to_track() 488 ) or ( 489 type(net.get_submodule(submodule_id)).__name__ 490 in GPA.pc.get_module_names_to_track() 491 ): 492 if GPA.pc.get_verbose(): 493 print( 494 "Seq sub is in tracking list so initiating tracked for: %s" 495 % sub_name 496 ) 497 if tracked_module_class is None: 498 continue 499 setattr( 500 net, 501 submodule_id, 502 tracked_module_class(net.get_submodule(submodule_id), sub_name), 503 ) 504 elif ( 505 type(net.get_submodule(submodule_id)) 506 in GPA.pc.get_modules_to_perforate() 507 or type(net.get_submodule(submodule_id)).__name__ 508 in GPA.pc.get_module_names_to_perforate() 509 ): 510 if GPA.pc.get_verbose(): 511 print( 512 "Seq sub is in conversion list so initing PAI for: " 513 "%s" % sub_name 514 ) 515 if ( 516 issubclass( 517 type(net.get_submodule(submodule_id)), 518 torch.nn.modules.batchnorm._BatchNorm, 519 ) 520 or issubclass( 521 type(net.get_submodule(submodule_id)), 522 torch.nn.modules.instancenorm._InstanceNorm, 523 ) 524 or issubclass( 525 type(net.get_submodule(submodule_id)), 526 torch.nn.modules.normalization.LayerNorm, 527 ) 528 ): 529 print( 530 "You have an unwrapped normalization layer, this " 531 "is not recommended: " + name_so_far 532 ) 533 pdb.set_trace() 534 setattr( 535 net, 536 submodule_id, 537 neuron_module_class(net.get_submodule(submodule_id), sub_name), 538 ) 539 else: 540 if net != net.get_submodule(submodule_id): 541 converted_list += [id(net.get_submodule(submodule_id))] 542 converted_names_list += [sub_name] 543 if GPA.pc.get_verbose(): 544 print( 545 "sub is module but in no lists so going deeper: %s" 546 % sub_name 547 ) 548 549 setattr( 550 net, 551 submodule_id, 552 convert_module( 553 net.get_submodule(submodule_id), 554 depth + 1, 555 sub_name, 556 converted_list, 557 converted_names_list, 558 neuron_module_class, 559 tracked_module_class, 560 ), 561 ) 562 # else: 563 # print('%s is a self pointer so skipping' % (name_so_far + '[' + str(submodule_id) + ']')) 564 elif type(net) in GPA.pc.get_modules_to_track(): 565 # print('skipping type for returning from call to: %s' % (name_so_far)) 566 return net 567 else: 568 for member in all_members: 569 if isinstance(getattr(type(net), member, None), property): 570 continue 571 # Immediately check if able to get the member, if not skip it 572 try: 573 getattr(net, member, None) 574 except: 575 continue 576 sub_name = name_so_far + "." + member 577 member_obj = getattr(net, member, None) 578 579 if isinstance(member_obj, (torch.nn.Parameter, torch.nn.parameter.Parameter)): 580 if sub_name in GPA.pc.get_parameter_ids_to_track(): 581 if GPA.pc.get_verbose(): 582 print("tracking parameter by ID: %s" % sub_name) 583 member_obj.parameter_type = "neuron" 584 member_obj.wrapped = True 585 continue 586 587 # Track module object ids once at this level so duplicate aliases are 588 # caught consistently (including direct children of the root module). 589 if isinstance(member_obj, nn.Module): 590 if id(member_obj) in converted_list: 591 original_sub_name = converted_names_list[ 592 converted_list.index(id(member_obj)) 593 ] 594 print( 595 "The following module has a duplicate pointer within " 596 "your model: %s" % sub_name 597 ) 598 print("Keeping first pointer: %s" % original_sub_name) 599 print("Skipping duplicate pointer: %s" % sub_name) 600 print( 601 "If you prefer to keep %s and skip %s, add %s to module_names_to_not_save before convert." 602 % (sub_name, original_sub_name, original_sub_name) 603 ) 604 GPA.pc.append_module_names_to_not_save([sub_name]) 605 continue 606 converted_list += [id(member_obj)] 607 converted_names_list += [sub_name] 608 if sub_name in GPA.pc.get_module_ids_to_track(): 609 if GPA.pc.get_verbose(): 610 print("Seq ID is in track IDs: %s" % sub_name) 611 if tracked_module_class is None: 612 continue 613 setattr( 614 net, member, tracked_module_class(getattr(net, member), sub_name) 615 ) 616 continue 617 if sub_name in GPA.pc.get_module_ids_to_perforate(): 618 if GPA.pc.get_verbose(): 619 print("Seq ID is in convert IDs: %s" % sub_name) 620 setattr( 621 net, member, neuron_module_class(getattr(net, member), sub_name) 622 ) 623 continue 624 if id(getattr(net, member, None)) == id(net): 625 if GPA.pc.get_verbose(): 626 print("member sub is a self pointer: %s" % sub_name) 627 continue 628 if sub_name in GPA.pc.get_module_names_to_not_save(): 629 if GPA.pc.get_verbose(): 630 print("Skipping %s during convert" % sub_name) 631 else: 632 if sub_name == ".base_model": 633 print( 634 "By default skipping base_model. See " 635 '"Safetensors Errors" section of ' 636 "customization.md to include it." 637 ) 638 continue 639 if type(getattr(net, member, None)) in GPA.pc.get_modules_to_replace(): 640 if GPA.pc.get_verbose(): 641 print("sub is in replacement module so replacing: %s" % sub_name) 642 setattr( 643 net, member, replace_predefined_modules(getattr(net, member, None)) 644 ) 645 if ( 646 type(getattr(net, member, None)) in GPA.pc.get_modules_to_track() 647 or type(getattr(net, member, None)).__name__ 648 in GPA.pc.get_module_names_to_track() 649 or sub_name in GPA.pc.get_module_ids_to_track() 650 ): 651 if GPA.pc.get_verbose(): 652 print( 653 "sub is in tracking list so initiating tracked for: %s" 654 % sub_name 655 ) 656 if tracked_module_class is None: 657 continue 658 setattr( 659 net, member, tracked_module_class(getattr(net, member), sub_name) 660 ) 661 elif ( 662 type(getattr(net, member, None)) in GPA.pc.get_modules_to_perforate() 663 or type(getattr(net, member, None)).__name__ 664 in GPA.pc.get_module_names_to_perforate() 665 or (sub_name in GPA.pc.get_module_ids_to_perforate()) 666 ): 667 if GPA.pc.get_verbose(): 668 print( 669 "sub is in conversion list so initiating PAI for: %s" % sub_name 670 ) 671 setattr( 672 net, 673 member, 674 neuron_module_class(getattr(net, member), sub_name), 675 ) 676 elif ( 677 issubclass(type(getattr(net, member, None)), nn.Module) 678 or issubclass(type(getattr(net, member, None)), nn.Sequential) 679 or issubclass(type(getattr(net, member, None)), nn.ModuleList) 680 ): 681 if net != getattr(net, member): 682 if GPA.pc.get_verbose(): 683 print( 684 "sub is module but in no lists so going deeper: %s" 685 % sub_name 686 ) 687 setattr( 688 net, 689 member, 690 convert_module( 691 getattr(net, member), 692 depth + 1, 693 sub_name, 694 converted_list, 695 converted_names_list, 696 neuron_module_class, 697 tracked_module_class, 698 ), 699 ) 700 if ( 701 issubclass( 702 type(getattr(net, member, None)), 703 torch.nn.modules.batchnorm._BatchNorm, 704 ) 705 or issubclass( 706 type(getattr(net, member, None)), 707 torch.nn.modules.instancenorm._InstanceNorm, 708 ) 709 or issubclass( 710 type(getattr(net, member, None)), 711 torch.nn.modules.normalization.LayerNorm, 712 ) 713 ): 714 if not GPA.pc.get_unwrapped_modules_confirmed(): 715 print( 716 "potentially found a norm Layer that " 717 "is not accounted for, this is not recommended: %s" % (sub_name) 718 ) 719 print( 720 "Set GPA.pc.set_unwrapped_modules_confirmed(True) to skip " 721 "this next time" 722 ) 723 print( 724 "inspect your network to " 725 "see what the module type containing this layer is." 726 ) 727 print("Then do one of the following:") 728 print( 729 " - Add the module type to " 730 "GPA.pc.get_module_names_to_perforate() to wrap it entirely" 731 ) 732 print( 733 " - If the norm layer is part of a sequential wrap " 734 "it and the previous layer in a PAISequential" 735 ) 736 print( 737 " - If you do not want to add dendrites to this " 738 "module add the type to GPA.pc.get_module_names_to_track()" 739 ) 740 pdb.set_trace() 741 else: 742 if GPA.pc.get_verbose(): 743 if member[0] != "_" or GPA.pc.get_extra_verbose() is True: 744 print("not calling convert on %s depth %d" % (member, depth)) 745 if GPA.pc.get_verbose(): 746 print("returning from call to: %s" % (name_so_far)) 747 return net 748 749 750def convert_network(net, layer_name=""): 751 """Function that calls convert_module and checks results 752 753 Parameters 754 ---------- 755 net : nn.Module 756 The network to convert. 757 layer_name : str, optional 758 The name of the layer if converting a single layer, by default "" 759 760 Returns 761 ------- 762 nn.Module 763 The converted network. 764 765 """ 766 if GPA.pc.get_perforated_backpropagation(): 767 UPB.initialize_pb() 768 MPB.set_main_parameters(net) 769 if type(net) in GPA.pc.get_modules_to_replace(): 770 net = replace_predefined_modules(net) 771 if (type(net) in GPA.pc.get_modules_to_perforate()) or ( 772 type(net).__name__ in GPA.pc.get_module_names_to_perforate() 773 ): 774 if layer_name == "": 775 print( 776 "converting a single layer without a name, add a " 777 "layer_name param to the call" 778 ) 779 sys.exit(-1) 780 net = PA.PAINeuronModule(net, layer_name) 781 else: 782 net = convert_module( 783 net, 0, "", [], [], PA.PAINeuronModule, PA.TrackedNeuronModule 784 ) 785 if GPA.pai_tracker.member_vars["doing_pai"]: 786 missed_ones = [] 787 tracked_ones = [] 788 for name, param in net.named_parameters(): 789 wrapped = "wrapped" in param.__dir__() 790 if wrapped: 791 if GPA.pc.get_verbose(): 792 print("param %s is now wrapped" % (name)) 793 else: 794 tracked = "tracked" in param.__dir__() 795 if tracked: 796 tracked_ones.append(name) 797 else: 798 missed_ones.append(name) 799 if ( 800 len(missed_ones) != 0 or len(tracked_ones) != 0 801 ) and GPA.pc.get_unwrapped_modules_confirmed() is False: 802 print( 803 "\n------------------------------------------------------------------" 804 ) 805 print( 806 "The following params are not wrapped.\n------------------------------------------------------------------" 807 ) 808 for name in tracked_ones: 809 print("." + name) 810 print( 811 "\n------------------------------------------------------------------" 812 ) 813 print( 814 "The following params are not tracked or wrapped.\n------------------------------------------------------------------" 815 ) 816 for name in missed_ones: 817 print("." + name) 818 print( 819 "\n------------------------------------------------------------------" 820 ) 821 print( 822 "Modules that are not wrapped will not have Dendrites to optimize them" 823 ) 824 print( 825 "Modules modules that are not tracked can cause errors and is NOT recommended" 826 ) 827 print( 828 "Any modules in the second list should be added to module_names_to_track" 829 ) 830 831 print( 832 "Set GPA.pc.set_unwrapped_modules_confirmed(True) to skip this next time" 833 ) 834 print( 835 "Inspect your network and see what the module types of these values are to add them to PGB.module_names_to_perforate" 836 ) 837 # If did miss some then set trace to debug 838 if len(missed_ones) != 0: 839 print( 840 "------------------------------------------------------------------\nType 'c' + enter to continue the run to confirm you do not want them to be refined" 841 ) 842 843 pdb.set_trace() 844 print("confirmed") 845 net.register_buffer("tracker_string", torch.tensor([], dtype=torch.uint8)) 846 return net 847 848 849def string_to_tensor(string): 850 """Helper function to convert a layer_tracker into a string 851 852 This is required for safetensors saving 853 854 Parameters 855 ---------- 856 string : str 857 The string to convert. 858 859 Returns 860 ------- 861 torch.Tensor 862 The converted tensor. 863 864 """ 865 ords = list(map(ord, string)) 866 ords = torch.tensor(ords, dtype=torch.uint8) 867 return ords 868 869 870def string_from_tensor(string_tensor): 871 """Convert a tensor back into a string 872 873 Parameters 874 ---------- 875 string_tensor : torch.Tensor 876 The tensor to convert. 877 878 Returns 879 ------- 880 str 881 The converted string. 882 883 """ 884 ords = string_tensor.tolist() 885 to_return = "" 886 # Doing block processing like this helps with memory errors 887 while len(ords) != 0: 888 remaining_ords = ords[100000:] 889 ords = ords[:100000] 890 to_append = "".join(map(chr, ords)) 891 to_return = to_return + to_append 892 ords = remaining_ords 893 return to_return 894 895 896def save_system(net, folder, name): 897 """Save the entire system 898 899 This saves the network itself as well as the tracker information 900 901 Parameters 902 ---------- 903 net : nn.Module 904 The network to save. 905 folder : str 906 The folder to save the network in. 907 name : str 908 The name to save the network under. 909 910 Returns 911 ------- 912 None 913 914 """ 915 if GPA.pc.get_verbose(): 916 print("saving system %s" % name) 917 temp = string_to_tensor(GPA.pai_tracker.to_string()) 918 if hasattr(net, "tracker_string"): 919 net.tracker_string = string_to_tensor(GPA.pai_tracker.to_string()).to( 920 next(net.parameters()).device 921 ) 922 else: 923 net.register_buffer( 924 "tracker_string", 925 string_to_tensor(GPA.pai_tracker.to_string()).to( 926 next(net.parameters()).device 927 ), 928 ) 929 # Before saving the tracker must be cleared to not contain pointers to the 930 # models modules 931 old_list = GPA.pai_tracker.neuron_module_vector 932 GPA.pai_tracker.neuron_module_vector = [] 933 save_net(net, folder, name) 934 GPA.pai_tracker.neuron_module_vector = old_list 935 pai_save_system(net, folder, name) 936 937 938def load_system( 939 net, 940 folder, 941 name, 942 load_from_restart=False, 943 switch_call=False, 944 load_from_manual_save=False, 945): 946 """Load the entire system 947 948 This is what should be used to load a saved system and restart training 949 950 Parameters 951 ---------- 952 net : nn.Module 953 The network to load into. 954 folder : str 955 The folder to load the network from. 956 name : str 957 The name to load the network from. 958 load_from_restart : bool, optional 959 Whether this is being loaded from an automatic restart, by default False 960 switch_call : bool, optional 961 Whether this is being called from a switch, by default False 962 load_from_manual_save : bool, optional 963 Whether this is being loaded from a manual save, by default False 964 965 Returns 966 ------- 967 nn.Module 968 The loaded network. 969 970 Notes 971 ----- 972 If you manually call save_system then load_from_manual_save should be True 973 974 """ 975 if GPA.pc.get_verbose(): 976 print("loading system %s" % name) 977 net = load_net(net, folder, name) 978 GPA.pai_tracker.reset_module_vector(net, load_from_restart) 979 980 GPA.pai_tracker.from_string(string_from_tensor(net.tracker_string)) 981 GPA.pai_tracker.saved_time = time.time() 982 GPA.pai_tracker.loaded = True 983 GPA.pai_tracker.member_vars["current_best_validation_score"] = 0 984 GPA.pai_tracker.member_vars["epoch_last_improved"] = GPA.pai_tracker.member_vars[ 985 "num_epochs_run" 986 ] 987 if GPA.pc.get_verbose(): 988 print( 989 "after loading epoch last improved is %d mode is %c" 990 % ( 991 GPA.pai_tracker.member_vars["epoch_last_improved"], 992 GPA.pai_tracker.member_vars["mode"], 993 ) 994 ) 995 996 # Saves always take place before the call to start_epoch so call it here 997 # when loading to correct off by 1 problems 998 if (not switch_call) and (not load_from_manual_save): 999 GPA.pai_tracker.start_epoch(internal_call=True) 1000 return net 1001 1002 1003def load_pretrained_model( 1004 net, 1005 folder, 1006 name, 1007 remove_dendrite_scaffolding=False, 1008): 1009 """Load a pretrained perforated model and reset tracker for fresh training. 1010 1011 This function loads a pretrained model's weights and dendrite structure while 1012 resetting all tracker state (epochs, switch history, etc.) to start training 1013 from scratch on a new task. This is useful for transfer learning where you want 1014 pretrained weights but need fresh training dynamics. 1015 1016 Parameters 1017 ---------- 1018 net : nn.Module 1019 The network to load into. 1020 folder : str 1021 The folder containing the pretrained model. 1022 name : str 1023 The name of the checkpoint to load (e.g., 'best_model', 'beforeSwitch_0'). 1024 remove_dendrite_scaffolding : bool, optional 1025 If True, removes dendrite scaffolding for inference or finetuning without 1026 adding more dendrites using blockwise_network and refresh_net. Default False. 1027 1028 Returns 1029 ------- 1030 nn.Module 1031 The loaded network with reset tracker state. 1032 1033 Examples 1034 -------- 1035 Load pretrained weights for continued dendrite training: 1036 >>> model = load_pretrained_model(model, "pretrained-prefc", "beforeSwitch_0") 1037 1038 Load pretrained weights for finetuning without adding more dendrites: 1039 >>> model = load_pretrained_model(model, "pretrained-prefc", "best_model", 1040 ... remove_dendrite_scaffolding=True) 1041 1042 Notes 1043 ----- 1044 This function: 1045 - Loads model weights and dendrite structure from checkpoint 1046 - Resets all epoch counters to -1 (will become 0 after first start_epoch) 1047 - Resets switch history and validation score tracking 1048 - Clears accuracy/loss history arrays 1049 - Optionally removes dendrite scaffolding (no more dendrite additions) 1050 1051 The tracker is reset to behave as if starting fresh training, while keeping 1052 the learned weights and dendrite structure from the pretrained model. 1053 """ 1054 from perforatedai import globals_perforatedai as GPA 1055 1056 if GPA.pc.get_verbose(): 1057 print(f"Loading pretrained model from {folder}/{name}") 1058 1059 # Load the model weights and dendrite structure 1060 net = load_system(net, folder, name, load_from_manual_save=True) 1061 1062 if GPA.pc.get_verbose(): 1063 print("Resetting tracker state for fresh training...") 1064 1065 # Reset structural training state to true initial values. 1066 # Keeping pretrained architecture/weights while zeroing cycle counters avoids 1067 # stale dendrite bookkeeping referencing empty score buffers. 1068 GPA.pai_tracker.reset_module_vector(net, load_from_restart=True) 1069 GPA.pai_tracker.member_vars["mode"] = "n" 1070 GPA.pai_tracker.member_vars["num_dendrites_added"] = 0 1071 GPA.pai_tracker.member_vars["num_dendrites_integrated"] = 0 1072 GPA.pai_tracker.member_vars["num_cycles"] = 0 1073 GPA.pai_tracker.member_vars["num_dendrite_tries"] = 0 1074 GPA.pai_tracker.member_vars["current_n_set_global_best"] = True 1075 1076 # Reset epoch counters 1077 GPA.pai_tracker.member_vars["num_epochs_run"] = -1 1078 GPA.pai_tracker.member_vars["total_epochs_run"] = -1 1079 GPA.pai_tracker.member_vars["epoch_last_improved"] = 0 1080 GPA.pai_tracker.member_vars["last_switch"] = 0 1081 GPA.pai_tracker.member_vars["manual_train_switch"] = False 1082 1083 # Reset switch history 1084 GPA.pai_tracker.member_vars["switch_epochs"] = [] 1085 GPA.pai_tracker.member_vars["n_switch_epochs"] = [] 1086 GPA.pai_tracker.member_vars["p_switch_epochs"] = [] 1087 GPA.pai_tracker.member_vars["param_counts"] = [] 1088 1089 # Reset validation scores and tracking 1090 GPA.pai_tracker.member_vars["current_best_validation_score"] = 0 1091 GPA.pai_tracker.member_vars["global_best_validation_score"] = 0 1092 GPA.pai_tracker.member_vars["running_accuracy"] = 0 1093 1094 # Clear accuracy/loss history arrays 1095 GPA.pai_tracker.member_vars["accuracies"] = [] 1096 GPA.pai_tracker.member_vars["last_improved_accuracies"] = [] 1097 GPA.pai_tracker.member_vars["test_accuracies"] = [] 1098 GPA.pai_tracker.member_vars["n_accuracies"] = [] 1099 GPA.pai_tracker.member_vars["p_accuracies"] = [] 1100 GPA.pai_tracker.member_vars["running_accuracies"] = [] 1101 GPA.pai_tracker.member_vars["training_loss"] = [] 1102 GPA.pai_tracker.member_vars["training_learning_rates"] = [] 1103 GPA.pai_tracker.member_vars["test_scores"] = [] 1104 1105 # Clear extra scores 1106 GPA.pai_tracker.member_vars["extra_scores"] = {} 1107 GPA.pai_tracker.member_vars["extra_scores_without_graphing"] = {} 1108 GPA.pai_tracker.member_vars["n_extra_scores"] = {} 1109 1110 # Keep per-layer dendrite score buffers initialized by reset_module_vector. 1111 1112 # Clear timing arrays 1113 GPA.pai_tracker.member_vars["n_epoch_times"] = [] 1114 GPA.pai_tracker.member_vars["p_epoch_times"] = [] 1115 GPA.pai_tracker.member_vars["n_train_times"] = [] 1116 GPA.pai_tracker.member_vars["p_train_times"] = [] 1117 GPA.pai_tracker.member_vars["n_val_times"] = [] 1118 GPA.pai_tracker.member_vars["p_val_times"] = [] 1119 1120 # Clear overwritten tracking 1121 GPA.pai_tracker.member_vars["overwritten_extras"] = [] 1122 GPA.pai_tracker.member_vars["overwritten_vals"] = [] 1123 GPA.pai_tracker.member_vars["overwritten_epochs"] = 0 1124 1125 # Reset learning rate search state 1126 GPA.pai_tracker.member_vars["initial_lr_test_epoch_count"] = -1 1127 GPA.pai_tracker.member_vars["current_n_learning_rate_initial_skip_steps"] = 0 1128 GPA.pai_tracker.member_vars["last_max_learning_rate_steps"] = 0 1129 GPA.pai_tracker.member_vars["last_max_learning_rate_value"] = -1 1130 GPA.pai_tracker.member_vars["current_cycle_lr_max_scores"] = [] 1131 GPA.pai_tracker.member_vars["current_step_count"] = 0 1132 GPA.pai_tracker.member_vars["committed_to_initial_rate"] = True 1133 GPA.pai_tracker.member_vars["best_mean_score_improved_this_epoch"] = 0 1134 GPA.pai_tracker.member_vars["step_status"] = TPA.STEP_CLEARED 1135 1136 # Reset saved time 1137 GPA.pai_tracker.start_time = time.time() 1138 GPA.pai_tracker.saved_time = 0 1139 1140 # Match tracker initialization behavior so first validation uses epoch 0. 1141 GPA.pai_tracker.start_epoch(internal_call=True) 1142 1143 if GPA.pc.get_verbose(): 1144 print( 1145 f"Tracker reset complete. Dendrites: {GPA.pai_tracker.member_vars['num_dendrites_integrated']}, " 1146 f"Mode: {GPA.pai_tracker.member_vars['mode']}" 1147 ) 1148 1149 # Optionally remove dendrite scaffolding 1150 if remove_dendrite_scaffolding: 1151 if GPA.pc.get_verbose(): 1152 print("Removing dendrite scaffolding (no dendrite additions)...") 1153 1154 from perforatedai import blockwise_perforatedai as BPA 1155 from perforatedai import clean_perforatedai as CPA 1156 1157 net = BPA.blockwise_network(net) 1158 net = CPA.refresh_net(net) 1159 1160 if GPA.pc.get_verbose(): 1161 print("Dendrite scaffolding removed. Model ready for inference or finetuning.") 1162 1163 return net 1164 1165 1166import json 1167from collections import defaultdict 1168from safetensors.torch import save_file, safe_open 1169import torch 1170 1171 1172def save_model_with_weight_tying(model, filepath): 1173 """Save model with safetensors while handling weight tying automatically 1174 1175 Parameters 1176 ---------- 1177 model : Any PyTorch Module. 1178 filepath : Path to filename. 1179 1180 Returns 1181 ------- 1182 dict[str, str] 1183 Mapping from secondary parameter keys to their primary tied key. 1184 """ 1185 state_dict = model.state_dict() 1186 1187 # Find all weight tied parameters 1188 tensor_to_keys = defaultdict(list) 1189 for key, tensor in state_dict.items(): 1190 # Use tensor data pointer as unique identifier 1191 tensor_id = tensor.data_ptr() 1192 tensor_to_keys[tensor_id].append(key) 1193 1194 # Find tied weights (tensors referenced by multiple keys) 1195 tied_weights = {} 1196 keys_to_remove = set() 1197 for tensor_id, keys in tensor_to_keys.items(): 1198 if len(keys) > 1 and not tensor_id == 0: 1199 # Multiple keys reference the same tensor - this is weight tying 1200 # Sort keys for deterministic ordering 1201 keys = sorted(keys) 1202 primary_key = keys[0] # Keep the first key 1203 for secondary_key in keys[1:]: 1204 tied_weights[secondary_key] = primary_key 1205 keys_to_remove.add(secondary_key) 1206 1207 # Remove tied weights from state_dict (keep only primary references) 1208 filtered_state_dict = { 1209 k: v for k, v in state_dict.items() if k not in keys_to_remove 1210 } 1211 1212 # Create metadata for weight tying information 1213 metadata = {} 1214 if tied_weights: 1215 # Store weight tying info as JSON string in metadata 1216 metadata["weight_tying"] = json.dumps(tied_weights) 1217 save_file(filtered_state_dict, filepath, metadata=metadata) 1218 print(f"Saved model with {len(tied_weights)} weight tying relationships") 1219 return tied_weights 1220 1221 1222def load_model_with_weight_tying(model, filepath): 1223 """Load model from safetensors while restoring weight tying 1224 1225 Parameters 1226 ---------- 1227 model : Any PyTorch module. 1228 filepath : Path to model file. 1229 1230 Returns 1231 ------- 1232 Loaded model. 1233 """ 1234 with safe_open(filepath, framework="pt") as f: 1235 metadata = f.metadata() 1236 state_dict = {key: f.get_tensor(key) for key in f.keys()} 1237 1238 # Restore weight tying if metadata exists 1239 tied_weights = {} 1240 if metadata and "weight_tying" in metadata: 1241 tied_weights = json.loads(metadata["weight_tying"]) 1242 for secondary_key, primary_key in tied_weights.items(): 1243 if primary_key in state_dict: 1244 # Restore the tied reference 1245 state_dict[secondary_key] = state_dict[primary_key] 1246 print(f"Restored weight tying: {secondary_key} -> {primary_key}") 1247 1248 # Handle tracker_string loading with flexible key matching 1249 tracker_key = None 1250 if "tracker_string" in state_dict: 1251 tracker_key = "tracker_string" 1252 else: 1253 # Search for keys containing "tracker_string" 1254 tracker_keys = [key for key in state_dict.keys() if "tracker_string" in key] 1255 if len(tracker_keys) == 1: 1256 tracker_key = tracker_keys[0] 1257 elif len(tracker_keys) > 1: 1258 print(f"Error: Multiple tracker_string keys found: {tracker_keys}") 1259 pdb.set_trace() 1260 else: 1261 print("Error: No tracker_string found in state_dict") 1262 1263 if tracker_key is not None and hasattr(model, "tracker_string"): 1264 model.tracker_string = state_dict[tracker_key] 1265 1266 model.load_state_dict(state_dict) 1267 return model 1268 1269 1270def save_net(net, folder, name): 1271 """Save the network 1272 1273 This is called within save_system after the tracker has been 1274 turned into a single tensor to be saved as a part of the network 1275 1276 Parameters 1277 ---------- 1278 net : nn.Module 1279 The network to save. 1280 folder : str 1281 The folder to save the network in. 1282 name : str 1283 The name to save the network under. 1284 1285 Returns 1286 ------- 1287 None 1288 1289 """ 1290 # If running a DDP only save with first thread 1291 if "RANK" in os.environ: 1292 if int(os.environ["RANK"]) != 0: 1293 return 1294 if not os.path.isdir(folder): 1295 os.makedirs(folder) 1296 save_point = folder + "/" 1297 if not os.path.isdir(save_point): 1298 os.mkdir(save_point) 1299 for param in net.parameters(): 1300 param.data = param.data.contiguous() 1301 if GPA.pc.get_using_safe_tensors(): 1302 if GPA.pc.get_weight_tying_experimental(): 1303 save_model_with_weight_tying(net, save_point + name + ".pt") 1304 else: 1305 # Strip the . so that the naming is the same for everywhere but it works with state_dict naming 1306 not_save = [ns.lstrip('.') for ns in GPA.pc.get_module_names_to_not_save()] 1307 state_dict = {k: v for k, v in net.state_dict().items() 1308 if not any(k.startswith(ns) for ns in not_save)} 1309 save_file(state_dict, save_point + name + ".pt") 1310 else: 1311 torch.save(net, save_point + name + ".pt") 1312 1313 1314def save_pai_net(net, folder, name): 1315 """Save the final pai network 1316 1317 This can be called after training to save the final network 1318 with all scaffolding removed so only the refined weights remain 1319 1320 Parameters 1321 ---------- 1322 net : nn.Module 1323 The network to save. 1324 folder : str 1325 The folder to save the network in. 1326 name : str 1327 The name to save the network under. 1328 1329 Returns 1330 ------- 1331 None 1332 1333 """ 1334 # if running a DDP only save with first thread 1335 if "RANK" in os.environ: 1336 if int(os.environ["RANK"]) != 0: 1337 return 1338 1339 # print('calling save: %s' % name) 1340 # GPA.pai_tracker.archive_layer() 1341 # These deep copys are required or the real model will also have its layers replaced 1342 net = prepare_final_model(net) 1343 if not os.path.isdir(folder): 1344 os.makedirs(folder) 1345 save_point = folder + "/" 1346 if not os.path.isdir(save_point): 1347 os.mkdir(save_point) 1348 1349 if GPA.pc.get_using_safe_tensors(): 1350 if GPA.pc.get_weight_tying_experimental(): 1351 save_model_with_weight_tying(net, save_point + name + "_pai.pt") 1352 else: 1353 save_file(net.state_dict(), save_point + name + "_pai.pt") 1354 else: 1355 torch.save(net, save_point + name + "_pai.pt") 1356 1357 1358def save_pai_net(net, folder, name): 1359 """Save the final pai network 1360 1361 This can be called after training to save the final network 1362 with all scaffolding removed so only the refined weights remain 1363 1364 Parameters 1365 ---------- 1366 net : nn.Module 1367 The network to save. 1368 folder : str 1369 The folder to save the network in. 1370 name : str 1371 The name to save the network under. 1372 1373 Returns 1374 ------- 1375 None 1376 1377 """ 1378 # if running a DDP only save with first thread 1379 if "RANK" in os.environ: 1380 if int(os.environ["RANK"]) != 0: 1381 return 1382 1383 # print('calling save: %s' % name) 1384 # GPA.pai_tracker.archive_layer() 1385 # These deep copys are required or the real model will also have its layers replaced 1386 net = prepare_final_model(net) 1387 if not os.path.isdir(folder): 1388 os.makedirs(folder) 1389 save_point = folder + "/" 1390 if not os.path.isdir(save_point): 1391 os.mkdir(save_point) 1392 1393 if GPA.pc.get_using_safe_tensors(): 1394 if GPA.pc.get_weight_tying_experimental(): 1395 save_model_with_weight_tying(net, save_point + name + "_pai.pt") 1396 else: 1397 save_file(net.state_dict(), save_point + name + "_pai.pt") 1398 else: 1399 torch.save(net, save_point + name + "_pai.pt") 1400 1401 1402def manual_load_state_dict(model, state_dict): 1403 """Load a state dict into a model key-by-key with relaxed checks. 1404 1405 Parameters 1406 ---------- 1407 model : nn.Module 1408 Model receiving parameters. 1409 state_dict : dict 1410 Source parameter dictionary. 1411 1412 Notes 1413 ----- 1414 Keys configured in ``module_names_to_not_save`` are skipped. 1415 1416 Returns 1417 ------- 1418 None 1419 This function does not return a value. 1420 """ 1421 own_state = model.state_dict() 1422 not_save = [ns.lstrip('.') for ns in GPA.pc.get_module_names_to_not_save()] 1423 for name, param in state_dict.items(): 1424 if any(name.startswith(ns) for ns in not_save): 1425 print("skipping loading %s based on module_names_to_not_save" % name) 1426 continue 1427 if name not in own_state: 1428 print(f"Warning: {name} not found in model state_dict") 1429 continue 1430 if isinstance(param, torch.nn.Parameter): 1431 # Backwards compatibility for serialized parameters 1432 param = param.data 1433 try: 1434 own_state[name].copy_(param) 1435 except Exception as e: 1436 print(f"Error loading {name}: {e}") 1437 print("Manual load complete") 1438 1439 1440def load_net(net, folder, name): 1441 """load the network 1442 1443 This is called within load_system after the tracker has been 1444 loaded 1445 1446 Parameters 1447 ---------- 1448 net : nn.Module 1449 The network to save. 1450 folder : str 1451 The folder to save the network in. 1452 name : str 1453 The name to save the network under. 1454 1455 Returns 1456 ------- 1457 nn.Module 1458 The loaded network. 1459 1460 """ 1461 save_point = folder + "/" 1462 if GPA.pc.get_using_safe_tensors(): 1463 model_path = save_point + name + ".pt" 1464 if GPA.pc.get_weight_tying_experimental(): 1465 return load_model_with_weight_tying(net, model_path) 1466 else: 1467 try: 1468 with safe_open(model_path, framework="pt") as f: 1469 metadata = f.metadata() 1470 if metadata and "weight_tying" in metadata: 1471 return load_model_with_weight_tying(net, model_path) 1472 except Exception: 1473 pass 1474 state_dict = load_file(model_path) 1475 else: 1476 # Different versions of torch require this change 1477 try: 1478 state_dict = torch.load( 1479 save_point + name + ".pt", 1480 map_location=torch.device("cpu"), 1481 weights_only=False, 1482 ).state_dict() 1483 except: 1484 try: 1485 state_dict = torch.load( 1486 save_point + name + ".pt", map_location=torch.device("cpu") 1487 ).state_dict() 1488 except: 1489 state_dict = torch.load( 1490 save_point + name + ".pt", map_location=torch.device("cpu") 1491 ) 1492 return load_net_from_dict(net, state_dict) 1493 1494 1495def get_module_base_name(module): 1496 """Normalize a wrapped module name for state-dict key lookup. 1497 1498 Parameters 1499 ---------- 1500 module : nn.Module 1501 Module containing a ``name`` attribute. 1502 1503 Returns 1504 ------- 1505 str 1506 Base name with leading dot and optional ``module.`` prefix removed. 1507 """ 1508 module_name = module.name 1509 # This should always be true 1510 if module_name[0] == ".": 1511 # strip "." 1512 module_name = module_name[1:] 1513 # If it was a dataparallel it will also have a module at the start 1514 # so strip that for loading 1515 if module_name[:6] == "module": 1516 module_name = module_name[7:] 1517 return module_name 1518 1519 1520def load_net_from_dict(net, state_dict): 1521 """load the network 1522 1523 This is called within load_net 1524 1525 Parameters 1526 ---------- 1527 net : nn.Module 1528 The network to save. 1529 state_dict : dict 1530 The state dictionary to load. 1531 1532 Returns 1533 ------- 1534 nn.Module 1535 The loaded network. 1536 1537 """ 1538 if GPA.pc.get_verbose(): 1539 print("loading net from dict") 1540 pai_modules = get_pai_modules(net, 0) 1541 if pai_modules == []: 1542 print( 1543 "PAI load_net and load_system uses a state_dict so it must be\n" 1544 "called with a net after perforate_model has been called" 1545 ) 1546 print( 1547 "This is being flagged because you are attempting to load a model\n" 1548 "that does not have any pai_modules in it. Confirm that you are calling\n" 1549 "perforate_model on the correct model, and the same model is the one\n" 1550 "being passed into add_validation_score" 1551 ) 1552 import pdb # This needs to be here for cython for some reason. 1553 pdb.set_trace() 1554 sys.exit(-1) 1555 if GPA.pc.get_verbose(): 1556 print( 1557 "setting up arrays and simulating cycles for %d pai modules" 1558 % len(pai_modules) 1559 ) 1560 not_save = GPA.pc.get_module_names_to_not_save() 1561 for module in pai_modules: 1562 if any(module.name.startswith(ns) for ns in not_save): 1563 print("skipping loading %s based on module_names_to_not_save" % module.name) 1564 continue 1565 # Set up name to be what will be saved in the state dict 1566 module_name = get_module_base_name(module) 1567 module.clear_dendrites() 1568 for tracker in module.dendrite_module.dendrite_values: 1569 try: 1570 tracker.setup_arrays( 1571 len( 1572 state_dict[ 1573 module_name + ".dendrite_module.dendrite_values.0.shape" 1574 ] 1575 ) 1576 ) 1577 except Exception as e: 1578 print(e) 1579 print( 1580 "This value is missing from the state dict\n" 1581 "When missing this value it typically means you\n" 1582 "converted a module but didn't actually use it in\n" 1583 "your forward and backward pass." 1584 ) 1585 print("module was: %s" % module.name) 1586 print("There are many reasons this can happen:") 1587 print( 1588 "\n1 - check your model definition and forward function and " 1589 "ensure this module is being used properly" 1590 ) 1591 print( 1592 "with GPA.pc.set_verbose(True) you can confirm this is the case if\n" 1593 'you do not see a "setting d shape for" this module at the first training batch.' 1594 ) 1595 print( 1596 "If this is the case, and it is correct to not be passing data through it\n" 1597 "Set it to be a tracked module with:\n" 1598 'GPA.pc.append_module_ids_to_track(["%s"]) to leave it out ' 1599 % module.name 1600 ) 1601 print( 1602 "\n2 - This can happen if you adjusted your model " 1603 "definition after calling perforate_model" 1604 ) 1605 print( 1606 "for example with torch.compile. If the module name " 1607 "printed above does not contain all modules leading " 1608 "to the main definition" 1609 ) 1610 print( 1611 "this is likely the case for your problem. Fix by " 1612 "calling perforate_model after all other model " 1613 "initialization steps" 1614 ) 1615 first_key = next(iter(state_dict.keys())) 1616 print( 1617 "\n3 - This can happen is if the model where you called perforate_model\n" 1618 "and the model within add_validation_score are not the same. \n" 1619 "Check if the module above and .%s have the same prefix\n" 1620 % first_key 1621 ) 1622 print( 1623 "if one starts with .model or .base etc and the other does not, this is the problem." 1624 ) 1625 1626 print( 1627 "\n4 - If you are using this module but then not actually including\n" 1628 "the correct output tensor in the forward. For example\n" 1629 "if you are using an LSTM and forwarding hidden instead of otput\n" 1630 "but your processors are set up to work with output" 1631 ) 1632 print( 1633 "\n5 - if you are not properly calling backward at all." 1634 " If this is the first module in your network it is more" 1635 "likely this is the problem." 1636 "One check in these cases is to make sure you do not call an initial validation score" 1637 "before the first backward call.\nIf you do this, while testing_dendrite_capacity is True" 1638 "this error will be triggered." 1639 ) 1640 print( 1641 "\n6 - You have converted a module that is in a frozen" 1642 " part of the network and thus no gradients are flowing" 1643 ) 1644 print( 1645 "\n7 - You are running multiple experiments at once with the same save_name." 1646 " When running concurrent trials be sure to add save_name=<unique_name> to perforate_model." 1647 ) 1648 import pdb # This needs to be here for cython for some reason. 1649 pdb.set_trace() 1650 1651 # Perform as many cycles as the state dict has 1652 num_cycles = int(state_dict[module_name + ".dendrite_module.num_cycles"].item()) 1653 if num_cycles > 0: 1654 simulate_cycles(module, num_cycles, doing_pai=True) 1655 # Handle tracker_string loading with flexible key matching 1656 tracker_key = None 1657 if "tracker_string" in state_dict: 1658 tracker_key = "tracker_string" 1659 else: 1660 # Search for keys containing "tracker_string" 1661 tracker_keys = [key for key in state_dict.keys() if "tracker_string" in key] 1662 if len(tracker_keys) == 1: 1663 tracker_key = tracker_keys[0] 1664 elif len(tracker_keys) > 1: 1665 print(f"Error: Multiple tracker_string keys found: {tracker_keys}") 1666 import pdb # This needs to be here for cython for some reason. 1667 pdb.set_trace() 1668 else: 1669 print("Error: No tracker_string found in state_dict") 1670 import pdb # This needs to be here for cython for some reason. 1671 pdb.set_trace() 1672 1673 if hasattr(net, "tracker_string"): 1674 net.tracker_string = state_dict[tracker_key] 1675 else: 1676 net.register_buffer("tracker_string", state_dict[tracker_key]) 1677 try: 1678 load_result = net.load_state_dict(state_dict, strict=False) 1679 not_save_state_names = [ns.lstrip('.') for ns in not_save] 1680 1681 def is_ignored_key(key): 1682 """Check whether a state-dict key should be ignored. 1683 1684 Parameters 1685 ---------- 1686 key : str 1687 State-dict key to test. 1688 1689 Returns 1690 ------- 1691 bool 1692 ``True`` when key belongs to a not-saved namespace. 1693 """ 1694 return any(key.startswith(ns) for ns in not_save_state_names) 1695 1696 missing_keys = [key for key in load_result.missing_keys if not is_ignored_key(key)] 1697 unexpected_keys = [key for key in load_result.unexpected_keys if not is_ignored_key(key)] 1698 1699 if GPA.pc.get_strict_loading() and (missing_keys or unexpected_keys): 1700 raise RuntimeError( 1701 "Error(s) in loading state_dict for %s:\n\tMissing key(s) in state_dict: %s. \n\tUnexpected key(s) in state_dict: %s." 1702 % (type(net).__name__, missing_keys, unexpected_keys) 1703 ) 1704 except Exception as e: 1705 """ 1706 When modules have high depth to them (i.e. modules within modules not number of layers) 1707 PyTorch can have trouble loading state dicts even when they are correct. 1708 This is a workaround to manually load the state dict if this happens. 1709 """ 1710 filtered_net_keys = { 1711 key 1712 for key in net.state_dict().keys() 1713 if not any(key.startswith(ns.lstrip('.')) for ns in not_save) 1714 } 1715 if filtered_net_keys == set(state_dict.keys()): 1716 print("Attempting manual loading of state_dict") 1717 manual_load_state_dict(net, state_dict) 1718 else: 1719 print(f"Error loading state_dict: {e}") 1720 print("If the error is due to missing keys (e.g., from code changes), you can try:") 1721 print(" GPA.pc.set_strict_loading(False)") 1722 print(" Do not change this unless you are certain the missing keys are not important to load and are expected due to code changes or arch changes.") 1723 print("\ntype 'c' to print full state dicts\n") 1724 import pdb # This needs to be here for cython for some reason. 1725 pdb.set_trace() 1726 print("net state dict is:") 1727 print(net.state_dict()) 1728 print("loaded state dict is:") 1729 print(state_dict) 1730 print( 1731 "Try to check differences. Likely is caused by a module not " 1732 "being converted that should be or vice versa" 1733 ) 1734 pdb.set_trace() 1735 net.to(GPA.pc.get_device()) 1736 return net 1737 1738 1739def pai_save_system(net, folder, name): 1740 """Save the entire system with scaffolding removed 1741 1742 This is used for the final network for inference after training 1743 1744 Parameters 1745 ---------- 1746 net : nn.Module 1747 The network to save. 1748 folder : str 1749 The folder to save the network in. 1750 name : str 1751 The name to save the network under. 1752 1753 Returns 1754 ------- 1755 None 1756 1757 """ 1758 net.member_vars = {} 1759 for member_var in GPA.pai_tracker.member_vars: 1760 if member_var == "scheduler_instance" or member_var == "optimizer_instance": 1761 continue 1762 net.member_vars[member_var] = GPA.pai_tracker.member_vars[member_var] 1763 pai_save_net(net, folder, name) 1764 1765 1766def deep_copy_pai(net): 1767 """Deep copy a PAI network 1768 1769 1770 Parameters 1771 ---------- 1772 net : nn.Module 1773 The network to copy. 1774 1775 Returns 1776 ------- 1777 nn.Module 1778 The copied network. 1779 1780 Notes 1781 ---- 1782 This is required because processors must be cleared before calling copy 1783 1784 """ 1785 # Dont check this stuff if its before the perforate_model has been called and you're just copying a regular model 1786 if(GPA.pai_tracker != []): 1787 # Clear gradients before saving the model 1788 if ((GPA.pai_tracker.member_vars["optimizer_instance"]) is not None) and ( 1789 GPA.pai_tracker.member_vars["optimizer_instance"] != [] 1790 ): 1791 GPA.pai_tracker.member_vars["optimizer_instance"].zero_grad() 1792 GPA.pai_tracker.clear_all_processors() 1793 return copy.deepcopy(net) 1794 1795 1796def prepare_final_model(net): 1797 """Prepare model for final save by removing scaffolding. 1798 1799 This performs all cleanup steps to convert a PAI model with scaffolding 1800 into a clean final model ready for inference or distribution. 1801 1802 Parameters 1803 ---------- 1804 net : nn.Module 1805 The network to prepare. 1806 1807 Returns 1808 ------- 1809 nn.Module 1810 The cleaned model with scaffolding removed. 1811 """ 1812 # Deep copy and clean the model (removes scaffolding) 1813 net = deep_copy_pai(net) 1814 net = BPA.blockwise_network(net) 1815 net = deep_copy_pai(net) 1816 net = CL.refresh_net(net) 1817 1818 # Remove tracker_string (not needed for final model) 1819 if hasattr(net, "tracker_string"): 1820 del net.tracker_string 1821 1822 # Make parameters contiguous 1823 for param in net.parameters(): 1824 param.data = param.data.contiguous() 1825 1826 return net 1827 1828 1829def pai_save_net(net, folder, name): 1830 """Save the entire system with scaffolding removed 1831 1832 This is called within pai_save_system after the tracker has been 1833 turned into a single tensor to be saved as a part of the network 1834 1835 1836 Parameters 1837 ---------- 1838 net : nn.Module 1839 The network to save. 1840 folder : str 1841 The folder to save the network in. 1842 name : str 1843 The name to save the network under. 1844 1845 Returns 1846 ------- 1847 None 1848 1849 Notes 1850 ---- 1851 For open source implementation this is not as important since 1852 minimal values are already being used. 1853 1854 """ 1855 1856 if GPA.pc.get_perforated_backpropagation(): 1857 UPB.pb_save_net(net, folder, name) 1858 else: 1859 return 1860 1861 1862def simulate_cycles(module, num_cycles, doing_pai): 1863 """Simulate dendrite addition cycles 1864 1865 Simulate the back and forth processes of adding dendrites to build a 1866 pretrained dendrite model before loading weights. Required for loading 1867 dendrite save files from non dendrite initial models. 1868 1869 Parameters 1870 ---------- 1871 module : PA.PAINeuronModule 1872 The module to simulate cycles on. 1873 num_cycles : int 1874 The number of cycles to simulate. 1875 doing_pai : bool 1876 Whether to actually do the simulation. 1877 1878 Returns 1879 ------- 1880 None 1881 1882 """ 1883 1884 check_skipped = GPA.pc.get_checked_skipped_modules() 1885 if doing_pai is False: 1886 return 1887 GPA.pc.set_checked_skipped_modules(True) 1888 mode = "n" 1889 for i in range(num_cycles): 1890 if mode == "n": 1891 module.set_mode("p") 1892 module.create_new_dendrite_module() 1893 mode = "p" 1894 else: 1895 module.set_mode("n") 1896 mode = "n" 1897 GPA.pc.set_checked_skipped_modules(check_skipped) 1898 1899 1900def count_params(net): 1901 """Count the number of parameters in the network 1902 1903 If doing perforated backpropagation this calls the PB function 1904 which does not count scaffolding parameters since the final model 1905 will not have them. 1906 1907 Parameters 1908 ---------- 1909 net : nn.Module 1910 The network to count parameters in. 1911 1912 Returns 1913 ------- 1914 int 1915 The number of parameters in the network. 1916 1917 """ 1918 if GPA.pc.get_perforated_backpropagation(): 1919 return UPB.pb_count_params(net) 1920 parameters = net.named_parameters() 1921 unique_params = { 1922 p.data_ptr(): p for name, p in parameters if "parent_module" not in name 1923 }.values() 1924 return sum(p.numel() for p in unique_params) 1925 1926 1927def change_learning_modes(net, folder, name, doing_pai): 1928 """Change between neuron and dendrite learning modes 1929 1930 High level steps for entire system to switch back and forth between 1931 neuron learning and dendrite learning 1932 1933 Parameters 1934 ---------- 1935 net : nn.Module 1936 The network to change modes on. 1937 folder : str 1938 The folder to save/load the network in/from. 1939 name : str 1940 The name to save/load the network under. 1941 doing_pai : bool 1942 Whether to add dendrites when changing modes. 1943 1944 Returns 1945 ------- 1946 int 1947 The number of parameters in the network. 1948 1949 Notes 1950 ----- 1951 If doing_pai is False this just allows training to continue longer rather than early stopping 1952 1953 """ 1954 # If not adding dendrites this just allows training to continue longer with flags 1955 # every time early stopping should be occurring 1956 if doing_pai is False: 1957 GPA.pai_tracker.member_vars["switch_epochs"].append( 1958 GPA.pai_tracker.member_vars["num_epochs_run"] 1959 ) 1960 GPA.pai_tracker.member_vars["last_switch"] = GPA.pai_tracker.member_vars[ 1961 "switch_epochs" 1962 ][-1] 1963 GPA.pai_tracker.reset_vals_for_score_reset() 1964 return net 1965 if GPA.pai_tracker.member_vars["mode"] == "n": 1966 current_epoch = GPA.pai_tracker.member_vars["num_epochs_run"] 1967 overwritten_epochs = GPA.pai_tracker.member_vars["overwritten_epochs"] 1968 overwritten_extra = GPA.pai_tracker.member_vars["extra_scores"] 1969 if GPA.pc.get_drawing_pai(): 1970 overwritten_val = GPA.pai_tracker.member_vars["accuracies"] 1971 else: 1972 overwritten_val = GPA.pai_tracker.member_vars["neuron_accuracies"] 1973 """ 1974 If true don't load the best system 1975 because it will delete dendrites if the previous best was better than 1976 the current best 1977 """ 1978 if not GPA.pc.get_silent(): 1979 print("Importing best Model for switch to PA...") 1980 net = load_system(net, folder, name, switch_call=True) 1981 GPA.pai_tracker.set_dendrite_training() 1982 GPA.pai_tracker.member_vars["overwritten_epochs"] = overwritten_epochs 1983 GPA.pai_tracker.member_vars["overwritten_epochs"] += ( 1984 current_epoch - GPA.pai_tracker.member_vars["num_epochs_run"] 1985 ) 1986 GPA.pai_tracker.member_vars["total_epochs_run"] = ( 1987 GPA.pai_tracker.member_vars["num_epochs_run"] 1988 + GPA.pai_tracker.member_vars["overwritten_epochs"] 1989 ) 1990 1991 if GPA.pc.get_save_old_graph_scores(): 1992 GPA.pai_tracker.member_vars["overwritten_extras"].append(overwritten_extra) 1993 GPA.pai_tracker.member_vars["overwritten_vals"].append(overwritten_val) 1994 else: 1995 GPA.pai_tracker.member_vars["overwritten_extras"] = [overwritten_extra] 1996 GPA.pai_tracker.member_vars["overwritten_vals"] = [overwritten_val] 1997 if GPA.pc.get_drawing_pai(): 1998 GPA.pai_tracker.member_vars["n_switch_epochs"].append( 1999 GPA.pai_tracker.member_vars["num_epochs_run"] 2000 ) 2001 else: 2002 if len(GPA.pai_tracker.member_vars["switch_epochs"]) == 0: 2003 GPA.pai_tracker.member_vars["n_switch_epochs"].append( 2004 GPA.pai_tracker.member_vars["num_epochs_run"] 2005 ) 2006 else: 2007 GPA.pai_tracker.member_vars["n_switch_epochs"].append( 2008 GPA.pai_tracker.member_vars["n_switch_epochs"][-1] 2009 + ( 2010 (GPA.pai_tracker.member_vars["num_epochs_run"]) 2011 - (GPA.pai_tracker.member_vars["switch_epochs"][-1]) 2012 ) 2013 ) 2014 2015 GPA.pai_tracker.member_vars["switch_epochs"].append( 2016 GPA.pai_tracker.member_vars["num_epochs_run"] 2017 ) 2018 GPA.pai_tracker.member_vars["last_switch"] = GPA.pai_tracker.member_vars[ 2019 "switch_epochs" 2020 ][-1] 2021 2022 # Because open source version is only doing neuron training for 2023 # gradient descent dendrites, switch back to n mode right away 2024 if ( 2025 not GPA.pc.get_perforated_backpropagation() 2026 ) or GPA.pc.get_no_extra_n_modes(): 2027 net = change_learning_modes(net, folder, name, doing_pai) 2028 else: 2029 if not GPA.pc.get_silent(): 2030 print("Switching back to N...") 2031 set_best = GPA.pai_tracker.member_vars["current_n_set_global_best"] 2032 GPA.pai_tracker.set_neuron_training() 2033 if len(GPA.pai_tracker.member_vars["p_switch_epochs"]) == 0: 2034 GPA.pai_tracker.member_vars["p_switch_epochs"].append( 2035 ( 2036 (GPA.pai_tracker.member_vars["num_epochs_run"] - 1) 2037 - (GPA.pai_tracker.member_vars["switch_epochs"][-1]) 2038 ) 2039 ) 2040 else: 2041 GPA.pai_tracker.member_vars["p_switch_epochs"].append( 2042 GPA.pai_tracker.member_vars["p_switch_epochs"][-1] 2043 + ( 2044 (GPA.pai_tracker.member_vars["num_epochs_run"]) 2045 - (GPA.pai_tracker.member_vars["switch_epochs"][-1]) 2046 ) 2047 ) 2048 GPA.pai_tracker.member_vars["switch_epochs"].append( 2049 GPA.pai_tracker.member_vars["num_epochs_run"] 2050 ) 2051 GPA.pai_tracker.member_vars["last_switch"] = GPA.pai_tracker.member_vars[ 2052 "switch_epochs" 2053 ][-1] 2054 # Will be false for open source implementation 2055 if GPA.pc.get_retain_all_dendrites() or ( 2056 GPA.pc.get_learn_dendrites_live() and set_best 2057 ): 2058 if not GPA.pc.get_silent(): 2059 print( 2060 "Saving model before starting normal training to " 2061 "retain PBNodes regardless of next N Phase results" 2062 ) 2063 save_system(net, folder, name) 2064 # if its just doing P for learn PAI live then switch back immediately 2065 if GPA.pc.get_perforated_backpropagation() and GPA.pc.get_no_extra_n_modes(): 2066 net = change_learning_modes(net, folder, name, doing_pai) 2067 2068 GPA.pai_tracker.member_vars["param_counts"].append(count_params(net)) 2069 2070 return net 2071 2072 2073def find_param_name_by_id(model, param_id): 2074 """ 2075 This is only used for debugging. 2076 Return the fully-qualified parameter name (e.g. "layer1.conv.weight") 2077 for the parameter whose id matches param_id. Returns None if not found. 2078 2079 This uses model.named_parameters(), which already recurses through submodules. 2080 2081 Parameters 2082 ---------- 2083 model : Model to look for param id. 2084 param_id : pointer to a parameter. 2085 2086 Returns 2087 ------- 2088 String representing the module within the model or None if not found. 2089 """ 2090 for name, p in model.named_parameters(recurse=True): 2091 if id(p) == param_id: 2092 return "." + name 2093 return None 2094 2095 2096def add_method_delegation_to_module(wrapper_module, method_name): 2097 """Add delegating methods to a wrapper module that has a main_module attribute. 2098 2099 This adds the specified methods to the wrapper module instance so they 2100 properly delegate to the wrapped main_module. Works for any wrapper module 2101 (TrackedNeuronModule, PAINeuronModule, etc.) that has a main_module attribute. 2102 2103 Args: 2104 wrapper_module: A wrapper module instance with a main_module attribute 2105 method_name: The method name to delegate (e.g., '_gradient_checkpointing_func') 2106 2107 Parameters 2108 ---------- 2109 wrapper_module : PyTorch Module that contains a sub module. 2110 method_name : Name of method to wrap. 2111 2112 Returns 2113 ------- 2114 None 2115 """ 2116 import types 2117 2118 if hasattr(wrapper_module.main_module, method_name): 2119 # Create a delegating method that forwards to main_module 2120 def make_delegated_method(name): 2121 """Create a bound delegation function for a given attribute name. 2122 2123 Parameters 2124 ---------- 2125 name : str 2126 Attribute name to forward to ``main_module``. 2127 2128 Returns 2129 ------- 2130 callable 2131 Function that delegates access or invocation. 2132 """ 2133 def delegated_method(self, *args, **kwargs): 2134 """Delegate attribute access or method call to ``main_module``. 2135 2136 Parameters 2137 ---------- 2138 *args : tuple 2139 Positional arguments forwarded to delegated callables. 2140 **kwargs : dict 2141 Keyword arguments forwarded to delegated callables. 2142 2143 Returns 2144 ------- 2145 Any 2146 Delegated attribute value or method result. 2147 """ 2148 main_module_attr = getattr(self.main_module, name, None) 2149 if main_module_attr is None: 2150 raise AttributeError( 2151 f"'{type(self.main_module).__name__}' object has no attribute '{name}'" 2152 ) 2153 if callable(main_module_attr): 2154 return main_module_attr(*args, **kwargs) 2155 return main_module_attr 2156 2157 return delegated_method 2158 2159 # Bind it to this specific instance 2160 setattr( 2161 wrapper_module, 2162 method_name, 2163 types.MethodType(make_delegated_method(method_name), wrapper_module), 2164 ) 2165 2166 2167def apply_method_delegation_to_model(model, method_name, main_module_type): 2168 """Recursively apply method delegation to all wrapper modules with main_module in a model. 2169 2170 This traverses the entire model and adds method delegation for any module that has 2171 a main_module attribute and optionally matches specified types. 2172 2173 Args: 2174 model: The PyTorch model to traverse 2175 method_name: The method name to delegate (e.g., '_gradient_checkpointing_func') 2176 main_module_type: main_module type name to filter by. 2177 Example: 'Qwen2DecoderLayer' 2178 2179 Example: 2180 # Apply gradient checkpointing delegation to all decoder layers 2181 apply_method_delegation_to_model( 2182 model, 2183 '_gradient_checkpointing_func', 2184 main_module_type='Qwen2DecoderLayer' 2185 ) 2186 2187 Parameters 2188 ---------- 2189 model : PyTorch model. 2190 method_name : method to delegate. 2191 main_module_type : type of module that has this method. 2192 2193 Returns 2194 ------- 2195 None 2196 This function does not return a value. 2197 """ 2198 count = 0 2199 for name, module in model.named_modules(): 2200 # Check if module has main_module attribute (it's a wrapper) 2201 if hasattr(module, "main_module"): 2202 # Check if we should apply based on main_module type 2203 should_apply = True 2204 if main_module_type is not None: 2205 main_module_type_name = type(module.main_module).__name__ 2206 should_apply = main_module_type_name == main_module_type 2207 2208 if should_apply: 2209 add_method_delegation_to_module(module, method_name) 2210 count += 1 2211 2212 print(f"[PAI] Applied method delegation to {count} wrapper module instances") 2213 2214 2215def make_json_serializable(obj): 2216 """Recursively convert non-JSON-serializable objects to strings. 2217 2218 Parameters 2219 ---------- 2220 obj : any 2221 The object to convert 2222 2223 Returns 2224 ------- 2225 Any 2226 JSON-serializable version of the object 2227 """ 2228 if isinstance(obj, (str, int, float, bool, type(None))): 2229 return obj 2230 elif isinstance(obj, dict): 2231 return {k: make_json_serializable(v) for k, v in obj.items()} 2232 elif isinstance(obj, (list, tuple)): 2233 return [make_json_serializable(item) for item in obj] 2234 else: 2235 # Convert non-serializable types to string 2236 return str(obj) 2237 2238 2239def extract_gpa_config(): 2240 """Extract all configuration from GPA.pc by calling all get_* methods. 2241 2242 Returns 2243 ------- 2244 dict[str, Any] 2245 Dictionary with all GPA.pc configuration values and type metadata 2246 2247 Examples 2248 -------- 2249 >>> config = extract_gpa_config() 2250 >>> # Returns: {'max_dendrites': 10, 'device': 'cuda', '_types': {...}} 2251 2252 Parameters 2253 ---------- 2254 None 2255 2256 """ 2257 config = {} 2258 config_types = {} 2259 2260 # Get all attributes from GPA.pc 2261 for attr_name in dir(GPA.pc): 2262 # Check if it starts with 'get_' 2263 if attr_name.startswith("get_"): 2264 try: 2265 # Get the method 2266 method = getattr(GPA.pc, attr_name) 2267 2268 # Check if it's callable 2269 if callable(method): 2270 # Call it and store result with key as name without 'get_' 2271 key = attr_name[4:] # Remove 'get_' prefix 2272 value = method() 2273 2274 # Check if this is an array (has corresponding append_ method) 2275 append_method_name = f"append_{key}" 2276 is_array = hasattr(GPA.pc, append_method_name) 2277 2278 if is_array and isinstance(value, (list, tuple)): 2279 # Store array element type 2280 if len(value) > 0: 2281 element_type = type(value[0]).__name__ 2282 else: 2283 element_type = None # empty array, no conversion needed 2284 config_types[key] = { 2285 "is_array": True, 2286 "element_type": element_type, 2287 } 2288 else: 2289 # Store value type 2290 config_types[key] = { 2291 "is_array": False, 2292 "type": type(value).__name__, 2293 } 2294 2295 # Make sure value is JSON serializable 2296 config[key] = make_json_serializable(value) 2297 except Exception as e: 2298 # Skip if method fails 2299 if GPA.pc.get_verbose(): 2300 print(f"Skipping {attr_name}: {e}") 2301 continue 2302 2303 # Add types metadata to config 2304 config["_types"] = config_types 2305 2306 return config 2307 2308 2309def convert_to_type(value, type_name): 2310 """Convert a value to the specified type. 2311 2312 Parameters 2313 ---------- 2314 value : any 2315 The value to convert 2316 type_name : str 2317 The target type name 2318 2319 Returns 2320 ------- 2321 Any 2322 The converted value 2323 """ 2324 if type_name == "NoneType" or value is None: 2325 return None 2326 elif type_name == "bool": 2327 if isinstance(value, str): 2328 return value.lower() in ("true", "1", "yes") 2329 return bool(value) 2330 elif type_name == "int": 2331 return int(value) 2332 elif type_name == "float": 2333 return float(value) 2334 elif type_name == "str": 2335 return str(value) 2336 elif type_name == "list": 2337 if not isinstance(value, list): 2338 return [value] 2339 return value 2340 elif type_name == "dict": 2341 if not isinstance(value, dict): 2342 return {} 2343 return value 2344 elif type_name == "type": 2345 # Handle type objects - convert string representation back to type 2346 if isinstance(value, str): 2347 # Try to evaluate the type string (e.g., "<class 'torch.nn.Linear'>") 2348 # Extract the class path from the string 2349 if value.startswith("<class '") and value.endswith("'>"): 2350 class_path = value[ 2351 8:-2 2352 ] # Extract 'torch.nn.Linear' from "<class 'torch.nn.Linear'>" 2353 parts = class_path.split(".") 2354 # Try to import and get the type 2355 try: 2356 module_name = ".".join(parts[:-1]) 2357 class_name = parts[-1] 2358 module = __import__(module_name, fromlist=[class_name]) 2359 return getattr(module, class_name) 2360 except Exception as e: 2361 print( 2362 f"Warning: Could not convert type string '{value}' to actual type: {e}" 2363 ) 2364 return value 2365 return value 2366 return value 2367 elif type_name == "dtype": 2368 # Handle torch dtype objects 2369 if isinstance(value, str): 2370 # Convert string like "torch.float32" to actual dtype 2371 import torch 2372 2373 try: 2374 # Try to get the dtype from torch module 2375 if value.startswith("torch."): 2376 dtype_name = value.split(".")[ 2377 1 2378 ] # Get 'float32' from 'torch.float32' 2379 return getattr(torch, dtype_name) 2380 else: 2381 return getattr(torch, value) 2382 except Exception as e: 2383 print( 2384 f"Warning: Could not convert dtype string '{value}' to actual dtype: {e}" 2385 ) 2386 return value 2387 return value 2388 elif type_name == "device": 2389 # Handle torch device objects 2390 if isinstance(value, str): 2391 # Convert string like "cuda" or "cpu" to torch.device 2392 import torch 2393 2394 try: 2395 return torch.device(value) 2396 except Exception as e: 2397 print( 2398 f"Warning: Could not convert device string '{value}' to actual device: {e}" 2399 ) 2400 return value 2401 return value 2402 elif type_name == "builtin_function_or_method": 2403 # Handle torch functions like torch.sigmoid, torch.relu, etc. 2404 if isinstance(value, str): 2405 # Parse string like "<built-in method sigmoid of type object at 0x...>" 2406 # to extract the function name 2407 import torch 2408 2409 try: 2410 if "<built-in method " in value and " of type object" in value: 2411 # Extract function name between '<built-in method ' and ' of type object' 2412 start = value.find("<built-in method ") + len("<built-in method ") 2413 end = value.find(" of type object") 2414 func_name = value[start:end] 2415 # Try to get the function from torch module 2416 if hasattr(torch, func_name): 2417 return getattr(torch, func_name) 2418 else: 2419 print(f"Warning: torch.{func_name} not found") 2420 return value 2421 else: 2422 return value 2423 except Exception as e: 2424 print( 2425 f"Warning: Could not convert builtin function string '{value}': {e}" 2426 ) 2427 return value 2428 return value 2429 else: 2430 # Unknown type - error and debug 2431 print(f"ERROR: Unknown type '{type_name}' for value: {value}") 2432 print(f"Type of value is: {type(value).__name__}") 2433 pdb.set_trace() 2434 return value 2435 2436 2437def convert_to_type_array(value, element_type): 2438 """Convert an array's elements to the specified type. 2439 2440 Parameters 2441 ---------- 2442 value : list or tuple 2443 The array to convert 2444 element_type : str or None 2445 The target type name for elements, None if array was empty 2446 2447 Returns 2448 ------- 2449 list 2450 The array with converted elements 2451 """ 2452 if not isinstance(value, (list, tuple)): 2453 return value 2454 # If element_type is None (empty array), no conversion needed 2455 if element_type is None: 2456 return list(value) if isinstance(value, tuple) else value 2457 return [convert_to_type(item, element_type) for item in value] 2458 2459 2460def set_gpa_config(config): 2461 """Set GPA.pc configuration by calling all set_* methods. 2462 2463 This is the reverse of extract_gpa_config(). It takes a configuration 2464 dictionary and calls the corresponding set_* methods on GPA.pc. 2465 Uses type metadata to ensure values are converted to the correct type. 2466 2467 Parameters 2468 ---------- 2469 config : dict 2470 Dictionary with configuration values (keys without 'set_' prefix) 2471 and optional '_types' metadata 2472 2473 Examples 2474 -------- 2475 >>> config = {'verbose': True, 'device': 'cuda'} 2476 >>> set_gpa_config(config) 2477 # Calls GPA.pc.set_verbose(True), GPA.pc.set_device('cuda'), etc. 2478 2479 Returns 2480 ------- 2481 Count of parameters that were set 2482 """ 2483 set_count = 0 2484 skip_count = 0 2485 2486 # Extract type information 2487 config_types = config.get("_types", {}) 2488 2489 for key, value in config.items(): 2490 # Skip the types metadata 2491 if key == "_types": 2492 continue 2493 2494 # Construct the set method name 2495 set_method_name = f"set_{key}" 2496 2497 # Check if the set method exists 2498 if hasattr(GPA.pc, set_method_name): 2499 try: 2500 method = getattr(GPA.pc, set_method_name) 2501 if callable(method): 2502 # Convert value to correct type if we have type info 2503 if key in config_types: 2504 type_info = config_types[key] 2505 if type_info.get("is_array", False): 2506 # Convert array elements to correct type 2507 element_type = type_info.get("element_type", "str") 2508 value = convert_to_type_array(value, element_type) 2509 else: 2510 # Convert single value to correct type 2511 value_type = type_info.get("type", "str") 2512 value = convert_to_type(value, value_type) 2513 2514 method(value) 2515 set_count += 1 2516 if GPA.pc.get_verbose(): 2517 print(f"Set {key} = {value}") 2518 except Exception as e: 2519 skip_count += 1 2520 if GPA.pc.get_verbose(): 2521 print(f"Failed to set {key}: {e}") 2522 else: 2523 skip_count += 1 2524 if GPA.pc.get_verbose(): 2525 print(f"No setter found for {key} (looking for {set_method_name})") 2526 2527 if GPA.pc.get_verbose(): 2528 print(f"Applied {set_count} PAI configuration settings ({skip_count} skipped)") 2529 2530 return set_count 2531 2532 2533try: 2534 from huggingface_hub import PyTorchModelHubMixin, hf_hub_download, HfApi 2535 2536 def upload_to_huggingface( 2537 model, 2538 repo_id, 2539 license="apache-2.0", 2540 pipeline_tag=None, 2541 repo_url=None, 2542 tags=None, 2543 include_pai_config=True, 2544 **kwargs, 2545 ): 2546 """Upload a model to HuggingFace Hub. 2547 2548 Uploads model weights and PAI configuration to HuggingFace Hub. 2549 The configuration is saved in config.json and can be restored when loading. 2550 2551 Parameters 2552 ---------- 2553 model : nn.Module 2554 The model to upload 2555 repo_id : str 2556 Repository ID (format: "username/model-name") 2557 license : str, optional 2558 License for the model card, by default "apache-2.0" 2559 pipeline_tag : str, optional 2560 Pipeline tag for the model (e.g., "text-classification", "image-classification") 2561 repo_url : str, optional 2562 URL to the model's repository/documentation 2563 tags : list, optional 2564 List of tags for the model card 2565 include_pai_config : bool, optional 2566 Whether to include all GPA.pc configuration in the model config, by default True 2567 **kwargs 2568 Additional arguments passed to HfApi (token, private, etc.) 2569 2570 Returns 2571 ------- 2572 str 2573 URL of the uploaded model 2574 2575 Examples 2576 -------- 2577 >>> url = upload_to_huggingface( 2578 ... model, 2579 ... "username/my-model", 2580 ... license="mit", 2581 ... pipeline_tag="image-classification", 2582 ... tags=["pytorch", "vision"] 2583 ... ) 2584 """ 2585 try: 2586 from huggingface_hub import HfApi 2587 except ImportError: 2588 raise ImportError( 2589 "huggingface_hub is required. Install it with: pip install huggingface_hub" 2590 ) 2591 2592 import tempfile 2593 import os 2594 2595 # Prepare model same way as save_pai_net does 2596 model = prepare_final_model(model) 2597 2598 # Calculate parameter count 2599 param_count = count_params(model) 2600 2601 # Format parameter count for tags (e.g., "11m" for 11 million) 2602 if param_count >= 1e9: 2603 param_tag = f"{param_count/1e9:.0f}b" 2604 elif param_count >= 1e6: 2605 param_tag = f"{param_count/1e6:.0f}m" 2606 elif param_count >= 1e3: 2607 param_tag = f"{param_count/1e3:.0f}k" 2608 else: 2609 param_tag = str(param_count) 2610 2611 # Create a temporary directory for files 2612 with tempfile.TemporaryDirectory() as tmpdir: 2613 # Save model weights 2614 model_path = os.path.join(tmpdir, "model.safetensors") 2615 save_file(model.state_dict(), model_path) 2616 2617 # Create config with PAI configuration 2618 config = {} 2619 if include_pai_config: 2620 pai_config = extract_gpa_config() 2621 config["pai_config"] = pai_config 2622 if GPA.pc.get_verbose(): 2623 print(f"Extracted {len(pai_config)} PAI configuration parameters") 2624 2625 # Add parameter count at top level 2626 config["num_parameters"] = param_count 2627 2628 # Add metadata 2629 if license: 2630 config["license"] = license 2631 if pipeline_tag: 2632 config["pipeline_tag"] = pipeline_tag 2633 if repo_url: 2634 config["repo_url"] = repo_url 2635 2636 # Add tags with parameter count 2637 if tags is None: 2638 tags = [] 2639 elif not isinstance(tags, list): 2640 tags = [tags] 2641 else: 2642 tags = tags.copy() # Don't modify the original list 2643 2644 # Add perforated-ai tag if not present 2645 if "perforated-ai" not in tags: 2646 tags.insert(0, "perforated-ai") 2647 2648 # Add parameter count tag if not present 2649 if param_tag not in tags: 2650 tags.append(param_tag) 2651 2652 config["tags"] = tags 2653 2654 # Save config.json 2655 config_path = os.path.join(tmpdir, "config.json") 2656 with open(config_path, "w") as f: 2657 json.dump(config, f, indent=2) 2658 2659 # Upload to HuggingFace 2660 api = HfApi() 2661 2662 # Extract token from kwargs if present 2663 token = kwargs.pop("token", None) 2664 private = kwargs.pop("private", None) 2665 2666 # Create repo if it doesn't exist 2667 try: 2668 api.create_repo( 2669 repo_id=repo_id, token=token, private=private, exist_ok=True 2670 ) 2671 except Exception as e: 2672 print(f"Repo may already exist: {e}") 2673 2674 # Upload folder 2675 api.upload_folder( 2676 folder_path=tmpdir, repo_id=repo_id, token=token, **kwargs 2677 ) 2678 2679 print(f"Model uploaded to: https://huggingface.co/{repo_id}") 2680 if include_pai_config: 2681 print(f"PAI configuration saved in config.json") 2682 print(f"To reload, use: model = from_hf_pretrained(model, '{repo_id}')") 2683 2684 return f"https://huggingface.co/{repo_id}" 2685 2686 def from_hf_pretrained(net, repo_id, force_download=False): 2687 """Load a PerforatedAI model from HuggingFace Hub using PyTorchModelHubMixin. 2688 2689 Args: 2690 net: The base model architecture (will be converted to PAI format) 2691 repo_id: HuggingFace Hub repository ID (e.g., "username/model-name") 2692 force_download: If True, always download the latest version, bypassing cache (default: False) 2693 2694 Returns: 2695 net: The loaded model with PAI modules initialized 2696 2697 Parameters 2698 ---------- 2699 net : PyTorch Model. 2700 repo_id : Name of HuggingFace repository. 2701 force_download : Force an update even if local file exists. 2702 2703 Returns 2704 ------- 2705 Loaded Model 2706 """ 2707 2708 # Wrap in a class that inherits from PyTorchModelHubMixin 2709 class PAIHFModel(net.__class__, PyTorchModelHubMixin): 2710 def __init__(self, *args, **kwargs): 2711 """Initialize the temporary HuggingFace-compatible wrapper.""" 2712 super().__init__(*args, **kwargs) 2713 2714 # Create an instance that can use from_pretrained 2715 wrapped_net = PAIHFModel.__new__(PAIHFModel) 2716 wrapped_net.__dict__ = net.__dict__ 2717 wrapped_net.__class__ = PAIHFModel 2718 2719 # Download config.json to restore PAI configuration 2720 try: 2721 config_path = hf_hub_download(repo_id=repo_id, filename="config.json", force_download=force_download) 2722 with open(config_path, "r") as f: 2723 config = json.load(f) 2724 if "pai_config" in config: 2725 # print(f"Restoring PAI configuration from HuggingFace") 2726 set_gpa_config(config["pai_config"]) 2727 else: 2728 print("Warning: No pai_config found in config.json") 2729 except Exception as e: 2730 print(f"Warning: Could not load PAI config from HuggingFace: {e}") 2731 2732 # Download model files from HuggingFace 2733 model_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", force_download=force_download) 2734 state_dict = load_file(model_path) 2735 wrapped_net = NPA.convert_network(wrapped_net) 2736 wrapped_net = NPA.load_pai_model_from_dict(wrapped_net, state_dict) 2737 return wrapped_net 2738 2739except: 2740 2741 def upload_to_huggingface(*args, **kwargs): 2742 """Raise an informative error when HuggingFace dependencies are missing. 2743 2744 Parameters 2745 ---------- 2746 *args : tuple[Any, ...] 2747 Positional arguments accepted for API compatibility. 2748 **kwargs : dict[str, Any] 2749 Keyword arguments accepted for API compatibility. 2750 2751 Returns 2752 ------- 2753 None 2754 Always raises ``ImportError``. 2755 """ 2756 raise ImportError( 2757 "huggingface_hub is required for upload_to_huggingface. " 2758 "Install it with: pip install huggingface_hub" 2759 ) 2760 2761 def from_hf_pretrained(*args, **kwargs): 2762 """Raise an informative error when HuggingFace dependencies are missing. 2763 2764 Parameters 2765 ---------- 2766 *args : tuple[Any, ...] 2767 Positional arguments accepted for API compatibility. 2768 **kwargs : dict[str, Any] 2769 Keyword arguments accepted for API compatibility. 2770 2771 Returns 2772 ------- 2773 None 2774 Always raises ``ImportError``. 2775 """ 2776 raise ImportError( 2777 "huggingface_hub is required for from_hf_pretrained. " 2778 "Install it with: pip install huggingface_hub" 2779 )
def
perforate_model( model, doing_pai=True, save_name='', making_graphs=True, maximizing_score=True, num_classes=10000000000, values_per_train_epoch=-1, values_per_val_epoch=-1, zooming_graph=True):
47def perforate_model( 48 model, 49 doing_pai=True, 50 save_name="", 51 making_graphs=True, 52 maximizing_score=True, 53 num_classes=10000000000, 54 values_per_train_epoch=-1, 55 values_per_val_epoch=-1, 56 zooming_graph=True, 57): 58 """Main function to initialize the network to add dendrites 59 60 This kicks off the entire Perforated AI process to add 61 the scaffolding to the network to be able to add dendrites 62 63 Parameters 64 ---------- 65 model : nn.Module 66 The neural network model to initialize. 67 doing_pai : bool, optional 68 Whether to actually add dendrites, by default True 69 save_name : str, optional 70 The name to save the model under, by default "PAI" 71 making_graphs : bool, optional 72 Whether to create graphs during training, by default True 73 maximizing_score : bool, optional 74 Whether to maximize the score during training, by default True 75 setting to false is for when the score is a loss to be minimized 76 num_classes : int, optional 77 The number of output classes, unused in current version 78 values_per_train_epoch : int, optional 79 The number of values to look back for graphing 80 during training, by default -1 (all values). 81 values_per_val_epoch : int, optional 82 The number of values to look back for graphing 83 during validation, by default -1 (all values). 84 zooming_graph : bool, optional 85 Whether to enable zooming on the graphs, by default True 86 87 Returns 88 ------- 89 model : nn.Module 90 The modified model with dendrite scaffolding added if doing_pai is True 91 92 """ 93 94 if save_name == "": 95 if GPA.pc.get_save_name() == "": 96 save_name = "PAI" 97 else: 98 save_name = GPA.pc.get_save_name() 99 100 if "/" in save_name: 101 print( 102 f"Warning: save_name '{save_name}' contains '/'. Relative paths are not implemented yet." 103 ) 104 sys.exit(1) 105 106 sanitized_save_name = "".join( 107 ch for ch in save_name if ch.isalnum() or ch in ("_", "-", ".") 108 ) 109 if sanitized_save_name != save_name: 110 print( 111 f"Warning: save_name '{save_name}' contained spaces or special characters. " 112 f"Using '{sanitized_save_name}' instead." 113 ) 114 save_name = sanitized_save_name 115 116 if save_name == "": 117 print("Warning: save_name became empty after sanitization. Using 'PAI'.") 118 save_name = "PAI" 119 120 121 GPA.pai_tracker = TPA.PAINeuronModuleTracker( 122 doing_pai=doing_pai, save_name=save_name 123 ) 124 GPA.pc.set_save_name(save_name) 125 if _dashboard_emitter is not None: 126 _dashboard_emitter.emit_run_start(GPA.pc, save_name) 127 model = GPA.pai_tracker.initialize( 128 model, 129 doing_pai=doing_pai, 130 save_name=save_name, 131 making_graphs=making_graphs, 132 maximizing_score=maximizing_score, 133 num_classes=num_classes, 134 values_per_train_epoch=-values_per_train_epoch, 135 values_per_val_epoch=values_per_val_epoch, 136 zooming_graph=zooming_graph, 137 ) 138 139 # Save config after perforation 140 if not GPA.pc.get_testing_dendrite_capacity(): 141 import os 142 GPA.pc.save_config(os.path.join(os.getcwd(), save_name, f"{save_name}_config.json")) 143 144 return model
Main function to initialize the network to add dendrites
This kicks off the entire Perforated AI process to add the scaffolding to the network to be able to add dendrites
Parameters
- model (nn.Module): The neural network model to initialize.
- doing_pai (bool, optional): Whether to actually add dendrites, by default True
- save_name (str, optional): The name to save the model under, by default "PAI"
- making_graphs (bool, optional): Whether to create graphs during training, by default True
- maximizing_score (bool, optional): Whether to maximize the score during training, by default True setting to false is for when the score is a loss to be minimized
- num_classes (int, optional): The number of output classes, unused in current version
- values_per_train_epoch (int, optional): The number of values to look back for graphing during training, by default -1 (all values).
- values_per_val_epoch (int, optional): The number of values to look back for graphing during validation, by default -1 (all values).
- zooming_graph (bool, optional): Whether to enable zooming on the graphs, by default True
Returns
- model (nn.Module): The modified model with dendrite scaffolding added if doing_pai is True
def
get_pai_modules(net, depth, seen_ids=None):
147def get_pai_modules(net, depth, seen_ids=None): 148 """Get a list of all neuron modules 149 150 Parameters 151 ---------- 152 net : nn.Module 153 The module to search. 154 depth : int 155 The current depth in the recursion. 156 157 Returns 158 ------- 159 list 160 A list of all PAI neuron modules found in the network. 161 162 """ 163 if seen_ids is None: 164 seen_ids = set() 165 all_members = net.__dir__() 166 this_list = [] 167 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 168 for submodule_id, layer in net.named_children(): 169 # If there is a self pointer ignore it 170 if net.get_submodule(submodule_id) is net: 171 continue 172 if type(net.get_submodule(submodule_id)) is PA.PAINeuronModule: 173 module = net.get_submodule(submodule_id) 174 if id(module) in seen_ids: 175 continue 176 seen_ids.add(id(module)) 177 this_list = this_list + [module] 178 else: 179 this_list = this_list + get_pai_modules( 180 net.get_submodule(submodule_id), depth + 1, seen_ids 181 ) 182 else: 183 for member in all_members: 184 if isinstance(getattr(type(net), member, None), property): 185 continue 186 # if the getter fails or it is a self pointer ignore it 187 try: 188 if getattr(net, member, None) is net: 189 continue 190 except: 191 continue 192 if type(getattr(net, member, None)) is PA.PAINeuronModule: 193 module = getattr(net, member) 194 if id(module) in seen_ids: 195 continue 196 seen_ids.add(id(module)) 197 this_list = this_list + [module] 198 elif ( 199 issubclass(type(getattr(net, member, None)), nn.Module) 200 or issubclass(type(getattr(net, member, None)), nn.Sequential) 201 or issubclass(type(getattr(net, member, None)), nn.ModuleList) 202 ): 203 this_list = this_list + get_pai_modules( 204 getattr(net, member), depth + 1, seen_ids 205 ) 206 207 return this_list
Get a list of all neuron modules
Parameters
- net (nn.Module): The module to search.
- depth (int): The current depth in the recursion.
Returns
- list: A list of all PAI neuron modules found in the network.
def
get_tracked_modules(net, depth, seen_ids=None):
210def get_tracked_modules(net, depth, seen_ids=None): 211 """Get a list of all tracked modules 212 213 Parameters 214 ---------- 215 net : nn.Module 216 The module to search. 217 depth : int 218 The current depth in the recursion. 219 220 Returns 221 ------- 222 list 223 A list of all tracked modules found in the network. 224 225 """ 226 if seen_ids is None: 227 seen_ids = set() 228 all_members = net.__dir__() 229 this_list = [] 230 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 231 for submodule_id, layer in net.named_children(): 232 if net.get_submodule(submodule_id) is net: 233 continue 234 if type(net.get_submodule(submodule_id)) is PA.TrackedNeuronModule: 235 module = net.get_submodule(submodule_id) 236 if id(module) in seen_ids: 237 continue 238 seen_ids.add(id(module)) 239 this_list = this_list + [module] 240 else: 241 this_list = this_list + get_tracked_modules( 242 net.get_submodule(submodule_id), depth + 1, seen_ids 243 ) 244 else: 245 for member in all_members: 246 if isinstance(getattr(type(net), member, None), property): 247 continue 248 # if the getter fails or it is a self pointer ignore it 249 try: 250 if getattr(net, member, None) is net: 251 continue 252 except: 253 continue 254 if type(getattr(net, member, None)) is PA.TrackedNeuronModule: 255 module = getattr(net, member) 256 if id(module) in seen_ids: 257 continue 258 seen_ids.add(id(module)) 259 this_list = this_list + [module] 260 elif issubclass(type(getattr(net, member, None)), nn.Module): 261 this_list = this_list + get_tracked_modules( 262 getattr(net, member), depth + 1, seen_ids 263 ) 264 return this_list
Get a list of all tracked modules
Parameters
- net (nn.Module): The module to search.
- depth (int): The current depth in the recursion.
Returns
- list: A list of all tracked modules found in the network.
def
get_pai_module_params(net, depth, seen_ids=None):
267def get_pai_module_params(net, depth, seen_ids=None): 268 """Get a list of all neuron module parameters 269 270 Parameters 271 ---------- 272 net : nn.Module 273 The module to search. 274 depth : int 275 The current depth in the recursion. 276 277 Returns 278 ------- 279 list 280 A list of all parameters of neuron modules found in this module. 281 282 """ 283 284 if seen_ids is None: 285 seen_ids = set() 286 all_members = net.__dir__() 287 this_list = [] 288 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 289 for submodule_id, layer in net.named_children(): 290 if isinstance(net.get_submodule(submodule_id), PA.PAINeuronModule): # 291 module = net.get_submodule(submodule_id) 292 if id(module) in seen_ids: 293 continue 294 seen_ids.add(id(module)) 295 for param in module.parameters(): 296 if param.requires_grad: 297 this_list = this_list + [param] 298 else: 299 this_list = this_list + get_pai_module_params( 300 net.get_submodule(submodule_id), depth + 1, seen_ids 301 ) 302 else: 303 for member in all_members: 304 if isinstance(getattr(type(net), member, None), property): 305 continue 306 if getattr(net, member, None) == net: 307 continue 308 if isinstance(getattr(net, member, None), PA.PAINeuronModule): 309 module = getattr(net, member) 310 if id(module) in seen_ids: 311 continue 312 seen_ids.add(id(module)) 313 for param in module.parameters(): 314 if param.requires_grad: 315 this_list = this_list + [param] 316 elif issubclass(type(getattr(net, member, None)), nn.Module): 317 this_list = this_list + get_pai_module_params( 318 getattr(net, member), depth + 1, seen_ids 319 ) 320 return this_list
Get a list of all neuron module parameters
Parameters
- net (nn.Module): The module to search.
- depth (int): The current depth in the recursion.
Returns
- list: A list of all parameters of neuron modules found in this module.
def
get_pai_network_params(net):
323def get_pai_network_params(net): 324 """Get a list of all neuron module parameters 325 326 Parameters 327 ---------- 328 net : nn.Module 329 The full model to search. 330 331 Returns 332 ------- 333 list 334 A list of all parameters of neuron modules found in the network. 335 336 """ 337 param_list = get_pai_module_params(net, 0) 338 return param_list
Get a list of all neuron module parameters
Parameters
- net (nn.Module): The full model to search.
Returns
- list: A list of all parameters of neuron modules found in the network.
def
replace_predefined_modules(start_module):
341def replace_predefined_modules(start_module): 342 """Replace a module with the module from globals list 343 344 Parameters 345 ---------- 346 start_module : nn.Module 347 The module to replace. 348 349 Returns 350 ------- 351 nn.Module 352 The replaced module. 353 354 """ 355 index = GPA.pc.get_modules_to_replace().index(type(start_module)) 356 return GPA.pc.get_replacement_modules()[index](start_module)
Replace a module with the module from globals list
Parameters
- start_module (nn.Module): The module to replace.
Returns
- nn.Module: The replaced module.
def
scan_module_aliases(net):
359def scan_module_aliases(net): 360 """Find alias module paths that point to already-seen module instances. 361 362 Parameters 363 ---------- 364 net : Any PyTorch Module. 365 366 Returns 367 ------- 368 dict[str, str] 369 Mapping of alias module paths to their canonical module paths. 370 """ 371 canonical = {} 372 aliases = {} 373 for name, module in net.named_modules(remove_duplicate=False): 374 if name == "": 375 continue 376 sub_name = "." + name 377 module_id = id(module) 378 if module_id in canonical: 379 aliases[sub_name] = canonical[module_id] 380 else: 381 canonical[module_id] = sub_name 382 return aliases
Find alias module paths that point to already-seen module instances.
Parameters
- net (Any PyTorch Module.):
Returns
- dict[str, str]: Mapping of alias module paths to their canonical module paths.
def
convert_module( net, depth, name_so_far, converted_list, converted_names_list, neuron_module_class, tracked_module_class):
385def convert_module( 386 net, 387 depth, 388 name_so_far, 389 converted_list, 390 converted_names_list, 391 neuron_module_class, 392 tracked_module_class, 393): 394 """Recursive function to do all conversion of modules to wrappers of modules 395 396 This is the function that goes through all of the module lists from 397 the globals file and does all the conversion and replacements to 398 setup the dendrite scaffolding as instructed. 399 400 Parameters 401 ---------- 402 net : nn.Module 403 The module to convert. 404 depth : int 405 The current depth in the recursion. 406 name_so_far : str 407 The name of the module so far in the recursion. 408 converted_list : list 409 A list of already converted module ids to avoid infinite loops. 410 converted_names_list : list 411 A corresponding list to help debug duplicate conversions 412 413 Returns 414 ------- 415 nn.Module 416 The converted module. 417 418 """ 419 if GPA.pc.get_verbose(): 420 print("calling convert on %s depth %d" % (net, depth)) 421 print( 422 "calling convert on %s: %s, depth %d" 423 % (name_so_far, type(net).__name__, depth) 424 ) 425 if isinstance(net, neuron_module_class) or ( 426 (tracked_module_class is not None) and isinstance(net, tracked_module_class) 427 ): 428 if GPA.pc.get_verbose(): 429 print( 430 "This is only being called because something in your model " 431 "is pointed to twice by two different variables. Highest " 432 "thing on the list is one of the duplicates" 433 ) 434 return net 435 if depth == 0 and name_so_far == "": 436 aliases = scan_module_aliases(net) 437 existing_not_save = set(GPA.pc.get_module_names_to_not_save()) 438 aliases_to_skip = [ 439 alias for alias in aliases.keys() if alias not in existing_not_save 440 ] 441 if aliases_to_skip: 442 GPA.pc.append_module_names_to_not_save(aliases_to_skip) 443 print( 444 "Auto-detected duplicate module aliases via named_modules; " 445 "keeping first-seen paths and skipping:" 446 ) 447 for alias in aliases_to_skip: 448 print(" - %s (keeps %s)" % (alias, aliases[alias])) 449 all_members = net.__dir__() 450 if GPA.pc.get_extra_verbose(): 451 print("all members:") 452 for member in all_members: 453 print(" - %s" % member) 454 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 455 for submodule_id, layer in net.named_children(): 456 sub_name = name_so_far + "." + str(submodule_id) 457 if sub_name in GPA.pc.get_module_ids_to_track(): 458 if GPA.pc.get_verbose(): 459 print("Seq ID is in track IDs: %s" % sub_name) 460 if tracked_module_class is None: 461 continue 462 setattr( 463 net, 464 submodule_id, 465 tracked_module_class(net.get_submodule(submodule_id), sub_name), 466 ) 467 continue 468 if sub_name in GPA.pc.get_module_ids_to_perforate(): 469 if GPA.pc.get_verbose(): 470 print("Seq ID is in convert IDs: %s" % sub_name) 471 setattr( 472 net, 473 submodule_id, 474 neuron_module_class(net.get_submodule(submodule_id), sub_name), 475 ) 476 continue 477 if type(net.get_submodule(submodule_id)) in GPA.pc.get_modules_to_replace(): 478 if GPA.pc.get_verbose(): 479 print( 480 "Seq sub is in replacement module so replacing: %s" % sub_name 481 ) 482 setattr( 483 net, 484 submodule_id, 485 replace_predefined_modules(net.get_submodule(submodule_id)), 486 ) 487 if ( 488 type(net.get_submodule(submodule_id)) in GPA.pc.get_modules_to_track() 489 ) or ( 490 type(net.get_submodule(submodule_id)).__name__ 491 in GPA.pc.get_module_names_to_track() 492 ): 493 if GPA.pc.get_verbose(): 494 print( 495 "Seq sub is in tracking list so initiating tracked for: %s" 496 % sub_name 497 ) 498 if tracked_module_class is None: 499 continue 500 setattr( 501 net, 502 submodule_id, 503 tracked_module_class(net.get_submodule(submodule_id), sub_name), 504 ) 505 elif ( 506 type(net.get_submodule(submodule_id)) 507 in GPA.pc.get_modules_to_perforate() 508 or type(net.get_submodule(submodule_id)).__name__ 509 in GPA.pc.get_module_names_to_perforate() 510 ): 511 if GPA.pc.get_verbose(): 512 print( 513 "Seq sub is in conversion list so initing PAI for: " 514 "%s" % sub_name 515 ) 516 if ( 517 issubclass( 518 type(net.get_submodule(submodule_id)), 519 torch.nn.modules.batchnorm._BatchNorm, 520 ) 521 or issubclass( 522 type(net.get_submodule(submodule_id)), 523 torch.nn.modules.instancenorm._InstanceNorm, 524 ) 525 or issubclass( 526 type(net.get_submodule(submodule_id)), 527 torch.nn.modules.normalization.LayerNorm, 528 ) 529 ): 530 print( 531 "You have an unwrapped normalization layer, this " 532 "is not recommended: " + name_so_far 533 ) 534 pdb.set_trace() 535 setattr( 536 net, 537 submodule_id, 538 neuron_module_class(net.get_submodule(submodule_id), sub_name), 539 ) 540 else: 541 if net != net.get_submodule(submodule_id): 542 converted_list += [id(net.get_submodule(submodule_id))] 543 converted_names_list += [sub_name] 544 if GPA.pc.get_verbose(): 545 print( 546 "sub is module but in no lists so going deeper: %s" 547 % sub_name 548 ) 549 550 setattr( 551 net, 552 submodule_id, 553 convert_module( 554 net.get_submodule(submodule_id), 555 depth + 1, 556 sub_name, 557 converted_list, 558 converted_names_list, 559 neuron_module_class, 560 tracked_module_class, 561 ), 562 ) 563 # else: 564 # print('%s is a self pointer so skipping' % (name_so_far + '[' + str(submodule_id) + ']')) 565 elif type(net) in GPA.pc.get_modules_to_track(): 566 # print('skipping type for returning from call to: %s' % (name_so_far)) 567 return net 568 else: 569 for member in all_members: 570 if isinstance(getattr(type(net), member, None), property): 571 continue 572 # Immediately check if able to get the member, if not skip it 573 try: 574 getattr(net, member, None) 575 except: 576 continue 577 sub_name = name_so_far + "." + member 578 member_obj = getattr(net, member, None) 579 580 if isinstance(member_obj, (torch.nn.Parameter, torch.nn.parameter.Parameter)): 581 if sub_name in GPA.pc.get_parameter_ids_to_track(): 582 if GPA.pc.get_verbose(): 583 print("tracking parameter by ID: %s" % sub_name) 584 member_obj.parameter_type = "neuron" 585 member_obj.wrapped = True 586 continue 587 588 # Track module object ids once at this level so duplicate aliases are 589 # caught consistently (including direct children of the root module). 590 if isinstance(member_obj, nn.Module): 591 if id(member_obj) in converted_list: 592 original_sub_name = converted_names_list[ 593 converted_list.index(id(member_obj)) 594 ] 595 print( 596 "The following module has a duplicate pointer within " 597 "your model: %s" % sub_name 598 ) 599 print("Keeping first pointer: %s" % original_sub_name) 600 print("Skipping duplicate pointer: %s" % sub_name) 601 print( 602 "If you prefer to keep %s and skip %s, add %s to module_names_to_not_save before convert." 603 % (sub_name, original_sub_name, original_sub_name) 604 ) 605 GPA.pc.append_module_names_to_not_save([sub_name]) 606 continue 607 converted_list += [id(member_obj)] 608 converted_names_list += [sub_name] 609 if sub_name in GPA.pc.get_module_ids_to_track(): 610 if GPA.pc.get_verbose(): 611 print("Seq ID is in track IDs: %s" % sub_name) 612 if tracked_module_class is None: 613 continue 614 setattr( 615 net, member, tracked_module_class(getattr(net, member), sub_name) 616 ) 617 continue 618 if sub_name in GPA.pc.get_module_ids_to_perforate(): 619 if GPA.pc.get_verbose(): 620 print("Seq ID is in convert IDs: %s" % sub_name) 621 setattr( 622 net, member, neuron_module_class(getattr(net, member), sub_name) 623 ) 624 continue 625 if id(getattr(net, member, None)) == id(net): 626 if GPA.pc.get_verbose(): 627 print("member sub is a self pointer: %s" % sub_name) 628 continue 629 if sub_name in GPA.pc.get_module_names_to_not_save(): 630 if GPA.pc.get_verbose(): 631 print("Skipping %s during convert" % sub_name) 632 else: 633 if sub_name == ".base_model": 634 print( 635 "By default skipping base_model. See " 636 '"Safetensors Errors" section of ' 637 "customization.md to include it." 638 ) 639 continue 640 if type(getattr(net, member, None)) in GPA.pc.get_modules_to_replace(): 641 if GPA.pc.get_verbose(): 642 print("sub is in replacement module so replacing: %s" % sub_name) 643 setattr( 644 net, member, replace_predefined_modules(getattr(net, member, None)) 645 ) 646 if ( 647 type(getattr(net, member, None)) in GPA.pc.get_modules_to_track() 648 or type(getattr(net, member, None)).__name__ 649 in GPA.pc.get_module_names_to_track() 650 or sub_name in GPA.pc.get_module_ids_to_track() 651 ): 652 if GPA.pc.get_verbose(): 653 print( 654 "sub is in tracking list so initiating tracked for: %s" 655 % sub_name 656 ) 657 if tracked_module_class is None: 658 continue 659 setattr( 660 net, member, tracked_module_class(getattr(net, member), sub_name) 661 ) 662 elif ( 663 type(getattr(net, member, None)) in GPA.pc.get_modules_to_perforate() 664 or type(getattr(net, member, None)).__name__ 665 in GPA.pc.get_module_names_to_perforate() 666 or (sub_name in GPA.pc.get_module_ids_to_perforate()) 667 ): 668 if GPA.pc.get_verbose(): 669 print( 670 "sub is in conversion list so initiating PAI for: %s" % sub_name 671 ) 672 setattr( 673 net, 674 member, 675 neuron_module_class(getattr(net, member), sub_name), 676 ) 677 elif ( 678 issubclass(type(getattr(net, member, None)), nn.Module) 679 or issubclass(type(getattr(net, member, None)), nn.Sequential) 680 or issubclass(type(getattr(net, member, None)), nn.ModuleList) 681 ): 682 if net != getattr(net, member): 683 if GPA.pc.get_verbose(): 684 print( 685 "sub is module but in no lists so going deeper: %s" 686 % sub_name 687 ) 688 setattr( 689 net, 690 member, 691 convert_module( 692 getattr(net, member), 693 depth + 1, 694 sub_name, 695 converted_list, 696 converted_names_list, 697 neuron_module_class, 698 tracked_module_class, 699 ), 700 ) 701 if ( 702 issubclass( 703 type(getattr(net, member, None)), 704 torch.nn.modules.batchnorm._BatchNorm, 705 ) 706 or issubclass( 707 type(getattr(net, member, None)), 708 torch.nn.modules.instancenorm._InstanceNorm, 709 ) 710 or issubclass( 711 type(getattr(net, member, None)), 712 torch.nn.modules.normalization.LayerNorm, 713 ) 714 ): 715 if not GPA.pc.get_unwrapped_modules_confirmed(): 716 print( 717 "potentially found a norm Layer that " 718 "is not accounted for, this is not recommended: %s" % (sub_name) 719 ) 720 print( 721 "Set GPA.pc.set_unwrapped_modules_confirmed(True) to skip " 722 "this next time" 723 ) 724 print( 725 "inspect your network to " 726 "see what the module type containing this layer is." 727 ) 728 print("Then do one of the following:") 729 print( 730 " - Add the module type to " 731 "GPA.pc.get_module_names_to_perforate() to wrap it entirely" 732 ) 733 print( 734 " - If the norm layer is part of a sequential wrap " 735 "it and the previous layer in a PAISequential" 736 ) 737 print( 738 " - If you do not want to add dendrites to this " 739 "module add the type to GPA.pc.get_module_names_to_track()" 740 ) 741 pdb.set_trace() 742 else: 743 if GPA.pc.get_verbose(): 744 if member[0] != "_" or GPA.pc.get_extra_verbose() is True: 745 print("not calling convert on %s depth %d" % (member, depth)) 746 if GPA.pc.get_verbose(): 747 print("returning from call to: %s" % (name_so_far)) 748 return net
Recursive function to do all conversion of modules to wrappers of modules
This is the function that goes through all of the module lists from the globals file and does all the conversion and replacements to setup the dendrite scaffolding as instructed.
Parameters
- net (nn.Module): The module to convert.
- depth (int): The current depth in the recursion.
- name_so_far (str): The name of the module so far in the recursion.
- converted_list (list): A list of already converted module ids to avoid infinite loops.
- converted_names_list (list): A corresponding list to help debug duplicate conversions
Returns
- nn.Module: The converted module.
def
convert_network(net, layer_name=''):
751def convert_network(net, layer_name=""): 752 """Function that calls convert_module and checks results 753 754 Parameters 755 ---------- 756 net : nn.Module 757 The network to convert. 758 layer_name : str, optional 759 The name of the layer if converting a single layer, by default "" 760 761 Returns 762 ------- 763 nn.Module 764 The converted network. 765 766 """ 767 if GPA.pc.get_perforated_backpropagation(): 768 UPB.initialize_pb() 769 MPB.set_main_parameters(net) 770 if type(net) in GPA.pc.get_modules_to_replace(): 771 net = replace_predefined_modules(net) 772 if (type(net) in GPA.pc.get_modules_to_perforate()) or ( 773 type(net).__name__ in GPA.pc.get_module_names_to_perforate() 774 ): 775 if layer_name == "": 776 print( 777 "converting a single layer without a name, add a " 778 "layer_name param to the call" 779 ) 780 sys.exit(-1) 781 net = PA.PAINeuronModule(net, layer_name) 782 else: 783 net = convert_module( 784 net, 0, "", [], [], PA.PAINeuronModule, PA.TrackedNeuronModule 785 ) 786 if GPA.pai_tracker.member_vars["doing_pai"]: 787 missed_ones = [] 788 tracked_ones = [] 789 for name, param in net.named_parameters(): 790 wrapped = "wrapped" in param.__dir__() 791 if wrapped: 792 if GPA.pc.get_verbose(): 793 print("param %s is now wrapped" % (name)) 794 else: 795 tracked = "tracked" in param.__dir__() 796 if tracked: 797 tracked_ones.append(name) 798 else: 799 missed_ones.append(name) 800 if ( 801 len(missed_ones) != 0 or len(tracked_ones) != 0 802 ) and GPA.pc.get_unwrapped_modules_confirmed() is False: 803 print( 804 "\n------------------------------------------------------------------" 805 ) 806 print( 807 "The following params are not wrapped.\n------------------------------------------------------------------" 808 ) 809 for name in tracked_ones: 810 print("." + name) 811 print( 812 "\n------------------------------------------------------------------" 813 ) 814 print( 815 "The following params are not tracked or wrapped.\n------------------------------------------------------------------" 816 ) 817 for name in missed_ones: 818 print("." + name) 819 print( 820 "\n------------------------------------------------------------------" 821 ) 822 print( 823 "Modules that are not wrapped will not have Dendrites to optimize them" 824 ) 825 print( 826 "Modules modules that are not tracked can cause errors and is NOT recommended" 827 ) 828 print( 829 "Any modules in the second list should be added to module_names_to_track" 830 ) 831 832 print( 833 "Set GPA.pc.set_unwrapped_modules_confirmed(True) to skip this next time" 834 ) 835 print( 836 "Inspect your network and see what the module types of these values are to add them to PGB.module_names_to_perforate" 837 ) 838 # If did miss some then set trace to debug 839 if len(missed_ones) != 0: 840 print( 841 "------------------------------------------------------------------\nType 'c' + enter to continue the run to confirm you do not want them to be refined" 842 ) 843 844 pdb.set_trace() 845 print("confirmed") 846 net.register_buffer("tracker_string", torch.tensor([], dtype=torch.uint8)) 847 return net
Function that calls convert_module and checks results
Parameters
- net (nn.Module): The network to convert.
- layer_name (str, optional): The name of the layer if converting a single layer, by default ""
Returns
- nn.Module: The converted network.
def
string_to_tensor(string):
850def string_to_tensor(string): 851 """Helper function to convert a layer_tracker into a string 852 853 This is required for safetensors saving 854 855 Parameters 856 ---------- 857 string : str 858 The string to convert. 859 860 Returns 861 ------- 862 torch.Tensor 863 The converted tensor. 864 865 """ 866 ords = list(map(ord, string)) 867 ords = torch.tensor(ords, dtype=torch.uint8) 868 return ords
Helper function to convert a layer_tracker into a string
This is required for safetensors saving
Parameters
- string (str): The string to convert.
Returns
- torch.Tensor: The converted tensor.
def
string_from_tensor(string_tensor):
871def string_from_tensor(string_tensor): 872 """Convert a tensor back into a string 873 874 Parameters 875 ---------- 876 string_tensor : torch.Tensor 877 The tensor to convert. 878 879 Returns 880 ------- 881 str 882 The converted string. 883 884 """ 885 ords = string_tensor.tolist() 886 to_return = "" 887 # Doing block processing like this helps with memory errors 888 while len(ords) != 0: 889 remaining_ords = ords[100000:] 890 ords = ords[:100000] 891 to_append = "".join(map(chr, ords)) 892 to_return = to_return + to_append 893 ords = remaining_ords 894 return to_return
Convert a tensor back into a string
Parameters
- string_tensor (torch.Tensor): The tensor to convert.
Returns
- str: The converted string.
def
save_system(net, folder, name):
897def save_system(net, folder, name): 898 """Save the entire system 899 900 This saves the network itself as well as the tracker information 901 902 Parameters 903 ---------- 904 net : nn.Module 905 The network to save. 906 folder : str 907 The folder to save the network in. 908 name : str 909 The name to save the network under. 910 911 Returns 912 ------- 913 None 914 915 """ 916 if GPA.pc.get_verbose(): 917 print("saving system %s" % name) 918 temp = string_to_tensor(GPA.pai_tracker.to_string()) 919 if hasattr(net, "tracker_string"): 920 net.tracker_string = string_to_tensor(GPA.pai_tracker.to_string()).to( 921 next(net.parameters()).device 922 ) 923 else: 924 net.register_buffer( 925 "tracker_string", 926 string_to_tensor(GPA.pai_tracker.to_string()).to( 927 next(net.parameters()).device 928 ), 929 ) 930 # Before saving the tracker must be cleared to not contain pointers to the 931 # models modules 932 old_list = GPA.pai_tracker.neuron_module_vector 933 GPA.pai_tracker.neuron_module_vector = [] 934 save_net(net, folder, name) 935 GPA.pai_tracker.neuron_module_vector = old_list 936 pai_save_system(net, folder, name)
Save the entire system
This saves the network itself as well as the tracker information
Parameters
- net (nn.Module): The network to save.
- folder (str): The folder to save the network in.
- name (str): The name to save the network under.
Returns
- None
def
load_system( net, folder, name, load_from_restart=False, switch_call=False, load_from_manual_save=False):
939def load_system( 940 net, 941 folder, 942 name, 943 load_from_restart=False, 944 switch_call=False, 945 load_from_manual_save=False, 946): 947 """Load the entire system 948 949 This is what should be used to load a saved system and restart training 950 951 Parameters 952 ---------- 953 net : nn.Module 954 The network to load into. 955 folder : str 956 The folder to load the network from. 957 name : str 958 The name to load the network from. 959 load_from_restart : bool, optional 960 Whether this is being loaded from an automatic restart, by default False 961 switch_call : bool, optional 962 Whether this is being called from a switch, by default False 963 load_from_manual_save : bool, optional 964 Whether this is being loaded from a manual save, by default False 965 966 Returns 967 ------- 968 nn.Module 969 The loaded network. 970 971 Notes 972 ----- 973 If you manually call save_system then load_from_manual_save should be True 974 975 """ 976 if GPA.pc.get_verbose(): 977 print("loading system %s" % name) 978 net = load_net(net, folder, name) 979 GPA.pai_tracker.reset_module_vector(net, load_from_restart) 980 981 GPA.pai_tracker.from_string(string_from_tensor(net.tracker_string)) 982 GPA.pai_tracker.saved_time = time.time() 983 GPA.pai_tracker.loaded = True 984 GPA.pai_tracker.member_vars["current_best_validation_score"] = 0 985 GPA.pai_tracker.member_vars["epoch_last_improved"] = GPA.pai_tracker.member_vars[ 986 "num_epochs_run" 987 ] 988 if GPA.pc.get_verbose(): 989 print( 990 "after loading epoch last improved is %d mode is %c" 991 % ( 992 GPA.pai_tracker.member_vars["epoch_last_improved"], 993 GPA.pai_tracker.member_vars["mode"], 994 ) 995 ) 996 997 # Saves always take place before the call to start_epoch so call it here 998 # when loading to correct off by 1 problems 999 if (not switch_call) and (not load_from_manual_save): 1000 GPA.pai_tracker.start_epoch(internal_call=True) 1001 return net
Load the entire system
This is what should be used to load a saved system and restart training
Parameters
- net (nn.Module): The network to load into.
- folder (str): The folder to load the network from.
- name (str): The name to load the network from.
- load_from_restart (bool, optional): Whether this is being loaded from an automatic restart, by default False
- switch_call (bool, optional): Whether this is being called from a switch, by default False
- load_from_manual_save (bool, optional): Whether this is being loaded from a manual save, by default False
Returns
- nn.Module: The loaded network.
Notes
If you manually call save_system then load_from_manual_save should be True
def
load_pretrained_model(net, folder, name, remove_dendrite_scaffolding=False):
1004def load_pretrained_model( 1005 net, 1006 folder, 1007 name, 1008 remove_dendrite_scaffolding=False, 1009): 1010 """Load a pretrained perforated model and reset tracker for fresh training. 1011 1012 This function loads a pretrained model's weights and dendrite structure while 1013 resetting all tracker state (epochs, switch history, etc.) to start training 1014 from scratch on a new task. This is useful for transfer learning where you want 1015 pretrained weights but need fresh training dynamics. 1016 1017 Parameters 1018 ---------- 1019 net : nn.Module 1020 The network to load into. 1021 folder : str 1022 The folder containing the pretrained model. 1023 name : str 1024 The name of the checkpoint to load (e.g., 'best_model', 'beforeSwitch_0'). 1025 remove_dendrite_scaffolding : bool, optional 1026 If True, removes dendrite scaffolding for inference or finetuning without 1027 adding more dendrites using blockwise_network and refresh_net. Default False. 1028 1029 Returns 1030 ------- 1031 nn.Module 1032 The loaded network with reset tracker state. 1033 1034 Examples 1035 -------- 1036 Load pretrained weights for continued dendrite training: 1037 >>> model = load_pretrained_model(model, "pretrained-prefc", "beforeSwitch_0") 1038 1039 Load pretrained weights for finetuning without adding more dendrites: 1040 >>> model = load_pretrained_model(model, "pretrained-prefc", "best_model", 1041 ... remove_dendrite_scaffolding=True) 1042 1043 Notes 1044 ----- 1045 This function: 1046 - Loads model weights and dendrite structure from checkpoint 1047 - Resets all epoch counters to -1 (will become 0 after first start_epoch) 1048 - Resets switch history and validation score tracking 1049 - Clears accuracy/loss history arrays 1050 - Optionally removes dendrite scaffolding (no more dendrite additions) 1051 1052 The tracker is reset to behave as if starting fresh training, while keeping 1053 the learned weights and dendrite structure from the pretrained model. 1054 """ 1055 from perforatedai import globals_perforatedai as GPA 1056 1057 if GPA.pc.get_verbose(): 1058 print(f"Loading pretrained model from {folder}/{name}") 1059 1060 # Load the model weights and dendrite structure 1061 net = load_system(net, folder, name, load_from_manual_save=True) 1062 1063 if GPA.pc.get_verbose(): 1064 print("Resetting tracker state for fresh training...") 1065 1066 # Reset structural training state to true initial values. 1067 # Keeping pretrained architecture/weights while zeroing cycle counters avoids 1068 # stale dendrite bookkeeping referencing empty score buffers. 1069 GPA.pai_tracker.reset_module_vector(net, load_from_restart=True) 1070 GPA.pai_tracker.member_vars["mode"] = "n" 1071 GPA.pai_tracker.member_vars["num_dendrites_added"] = 0 1072 GPA.pai_tracker.member_vars["num_dendrites_integrated"] = 0 1073 GPA.pai_tracker.member_vars["num_cycles"] = 0 1074 GPA.pai_tracker.member_vars["num_dendrite_tries"] = 0 1075 GPA.pai_tracker.member_vars["current_n_set_global_best"] = True 1076 1077 # Reset epoch counters 1078 GPA.pai_tracker.member_vars["num_epochs_run"] = -1 1079 GPA.pai_tracker.member_vars["total_epochs_run"] = -1 1080 GPA.pai_tracker.member_vars["epoch_last_improved"] = 0 1081 GPA.pai_tracker.member_vars["last_switch"] = 0 1082 GPA.pai_tracker.member_vars["manual_train_switch"] = False 1083 1084 # Reset switch history 1085 GPA.pai_tracker.member_vars["switch_epochs"] = [] 1086 GPA.pai_tracker.member_vars["n_switch_epochs"] = [] 1087 GPA.pai_tracker.member_vars["p_switch_epochs"] = [] 1088 GPA.pai_tracker.member_vars["param_counts"] = [] 1089 1090 # Reset validation scores and tracking 1091 GPA.pai_tracker.member_vars["current_best_validation_score"] = 0 1092 GPA.pai_tracker.member_vars["global_best_validation_score"] = 0 1093 GPA.pai_tracker.member_vars["running_accuracy"] = 0 1094 1095 # Clear accuracy/loss history arrays 1096 GPA.pai_tracker.member_vars["accuracies"] = [] 1097 GPA.pai_tracker.member_vars["last_improved_accuracies"] = [] 1098 GPA.pai_tracker.member_vars["test_accuracies"] = [] 1099 GPA.pai_tracker.member_vars["n_accuracies"] = [] 1100 GPA.pai_tracker.member_vars["p_accuracies"] = [] 1101 GPA.pai_tracker.member_vars["running_accuracies"] = [] 1102 GPA.pai_tracker.member_vars["training_loss"] = [] 1103 GPA.pai_tracker.member_vars["training_learning_rates"] = [] 1104 GPA.pai_tracker.member_vars["test_scores"] = [] 1105 1106 # Clear extra scores 1107 GPA.pai_tracker.member_vars["extra_scores"] = {} 1108 GPA.pai_tracker.member_vars["extra_scores_without_graphing"] = {} 1109 GPA.pai_tracker.member_vars["n_extra_scores"] = {} 1110 1111 # Keep per-layer dendrite score buffers initialized by reset_module_vector. 1112 1113 # Clear timing arrays 1114 GPA.pai_tracker.member_vars["n_epoch_times"] = [] 1115 GPA.pai_tracker.member_vars["p_epoch_times"] = [] 1116 GPA.pai_tracker.member_vars["n_train_times"] = [] 1117 GPA.pai_tracker.member_vars["p_train_times"] = [] 1118 GPA.pai_tracker.member_vars["n_val_times"] = [] 1119 GPA.pai_tracker.member_vars["p_val_times"] = [] 1120 1121 # Clear overwritten tracking 1122 GPA.pai_tracker.member_vars["overwritten_extras"] = [] 1123 GPA.pai_tracker.member_vars["overwritten_vals"] = [] 1124 GPA.pai_tracker.member_vars["overwritten_epochs"] = 0 1125 1126 # Reset learning rate search state 1127 GPA.pai_tracker.member_vars["initial_lr_test_epoch_count"] = -1 1128 GPA.pai_tracker.member_vars["current_n_learning_rate_initial_skip_steps"] = 0 1129 GPA.pai_tracker.member_vars["last_max_learning_rate_steps"] = 0 1130 GPA.pai_tracker.member_vars["last_max_learning_rate_value"] = -1 1131 GPA.pai_tracker.member_vars["current_cycle_lr_max_scores"] = [] 1132 GPA.pai_tracker.member_vars["current_step_count"] = 0 1133 GPA.pai_tracker.member_vars["committed_to_initial_rate"] = True 1134 GPA.pai_tracker.member_vars["best_mean_score_improved_this_epoch"] = 0 1135 GPA.pai_tracker.member_vars["step_status"] = TPA.STEP_CLEARED 1136 1137 # Reset saved time 1138 GPA.pai_tracker.start_time = time.time() 1139 GPA.pai_tracker.saved_time = 0 1140 1141 # Match tracker initialization behavior so first validation uses epoch 0. 1142 GPA.pai_tracker.start_epoch(internal_call=True) 1143 1144 if GPA.pc.get_verbose(): 1145 print( 1146 f"Tracker reset complete. Dendrites: {GPA.pai_tracker.member_vars['num_dendrites_integrated']}, " 1147 f"Mode: {GPA.pai_tracker.member_vars['mode']}" 1148 ) 1149 1150 # Optionally remove dendrite scaffolding 1151 if remove_dendrite_scaffolding: 1152 if GPA.pc.get_verbose(): 1153 print("Removing dendrite scaffolding (no dendrite additions)...") 1154 1155 from perforatedai import blockwise_perforatedai as BPA 1156 from perforatedai import clean_perforatedai as CPA 1157 1158 net = BPA.blockwise_network(net) 1159 net = CPA.refresh_net(net) 1160 1161 if GPA.pc.get_verbose(): 1162 print("Dendrite scaffolding removed. Model ready for inference or finetuning.") 1163 1164 return net
Load a pretrained perforated model and reset tracker for fresh training.
This function loads a pretrained model's weights and dendrite structure while resetting all tracker state (epochs, switch history, etc.) to start training from scratch on a new task. This is useful for transfer learning where you want pretrained weights but need fresh training dynamics.
Parameters
- net (nn.Module): The network to load into.
- folder (str): The folder containing the pretrained model.
- name (str): The name of the checkpoint to load (e.g., 'best_model', 'beforeSwitch_0').
- remove_dendrite_scaffolding (bool, optional): If True, removes dendrite scaffolding for inference or finetuning without adding more dendrites using blockwise_network and refresh_net. Default False.
Returns
- nn.Module: The loaded network with reset tracker state.
Examples
Load pretrained weights for continued dendrite training:
>>> model = load_pretrained_model(model, "pretrained-prefc", "beforeSwitch_0")
Load pretrained weights for finetuning without adding more dendrites:
>>> model = load_pretrained_model(model, "pretrained-prefc", "best_model",
... remove_dendrite_scaffolding=True)
Notes
This function:
- Loads model weights and dendrite structure from checkpoint
- Resets all epoch counters to -1 (will become 0 after first start_epoch)
- Resets switch history and validation score tracking
- Clears accuracy/loss history arrays
- Optionally removes dendrite scaffolding (no more dendrite additions)
The tracker is reset to behave as if starting fresh training, while keeping the learned weights and dendrite structure from the pretrained model.
def
save_model_with_weight_tying(model, filepath):
1173def save_model_with_weight_tying(model, filepath): 1174 """Save model with safetensors while handling weight tying automatically 1175 1176 Parameters 1177 ---------- 1178 model : Any PyTorch Module. 1179 filepath : Path to filename. 1180 1181 Returns 1182 ------- 1183 dict[str, str] 1184 Mapping from secondary parameter keys to their primary tied key. 1185 """ 1186 state_dict = model.state_dict() 1187 1188 # Find all weight tied parameters 1189 tensor_to_keys = defaultdict(list) 1190 for key, tensor in state_dict.items(): 1191 # Use tensor data pointer as unique identifier 1192 tensor_id = tensor.data_ptr() 1193 tensor_to_keys[tensor_id].append(key) 1194 1195 # Find tied weights (tensors referenced by multiple keys) 1196 tied_weights = {} 1197 keys_to_remove = set() 1198 for tensor_id, keys in tensor_to_keys.items(): 1199 if len(keys) > 1 and not tensor_id == 0: 1200 # Multiple keys reference the same tensor - this is weight tying 1201 # Sort keys for deterministic ordering 1202 keys = sorted(keys) 1203 primary_key = keys[0] # Keep the first key 1204 for secondary_key in keys[1:]: 1205 tied_weights[secondary_key] = primary_key 1206 keys_to_remove.add(secondary_key) 1207 1208 # Remove tied weights from state_dict (keep only primary references) 1209 filtered_state_dict = { 1210 k: v for k, v in state_dict.items() if k not in keys_to_remove 1211 } 1212 1213 # Create metadata for weight tying information 1214 metadata = {} 1215 if tied_weights: 1216 # Store weight tying info as JSON string in metadata 1217 metadata["weight_tying"] = json.dumps(tied_weights) 1218 save_file(filtered_state_dict, filepath, metadata=metadata) 1219 print(f"Saved model with {len(tied_weights)} weight tying relationships") 1220 return tied_weights
Save model with safetensors while handling weight tying automatically
Parameters
model (Any PyTorch Module.):
filepath (Path to filename.):
Returns
- dict[str, str]: Mapping from secondary parameter keys to their primary tied key.
def
load_model_with_weight_tying(model, filepath):
1223def load_model_with_weight_tying(model, filepath): 1224 """Load model from safetensors while restoring weight tying 1225 1226 Parameters 1227 ---------- 1228 model : Any PyTorch module. 1229 filepath : Path to model file. 1230 1231 Returns 1232 ------- 1233 Loaded model. 1234 """ 1235 with safe_open(filepath, framework="pt") as f: 1236 metadata = f.metadata() 1237 state_dict = {key: f.get_tensor(key) for key in f.keys()} 1238 1239 # Restore weight tying if metadata exists 1240 tied_weights = {} 1241 if metadata and "weight_tying" in metadata: 1242 tied_weights = json.loads(metadata["weight_tying"]) 1243 for secondary_key, primary_key in tied_weights.items(): 1244 if primary_key in state_dict: 1245 # Restore the tied reference 1246 state_dict[secondary_key] = state_dict[primary_key] 1247 print(f"Restored weight tying: {secondary_key} -> {primary_key}") 1248 1249 # Handle tracker_string loading with flexible key matching 1250 tracker_key = None 1251 if "tracker_string" in state_dict: 1252 tracker_key = "tracker_string" 1253 else: 1254 # Search for keys containing "tracker_string" 1255 tracker_keys = [key for key in state_dict.keys() if "tracker_string" in key] 1256 if len(tracker_keys) == 1: 1257 tracker_key = tracker_keys[0] 1258 elif len(tracker_keys) > 1: 1259 print(f"Error: Multiple tracker_string keys found: {tracker_keys}") 1260 pdb.set_trace() 1261 else: 1262 print("Error: No tracker_string found in state_dict") 1263 1264 if tracker_key is not None and hasattr(model, "tracker_string"): 1265 model.tracker_string = state_dict[tracker_key] 1266 1267 model.load_state_dict(state_dict) 1268 return model
Load model from safetensors while restoring weight tying
Parameters
model (Any PyTorch module.):
filepath (Path to model file.):
Returns
- Loaded model.
def
save_net(net, folder, name):
1271def save_net(net, folder, name): 1272 """Save the network 1273 1274 This is called within save_system after the tracker has been 1275 turned into a single tensor to be saved as a part of the network 1276 1277 Parameters 1278 ---------- 1279 net : nn.Module 1280 The network to save. 1281 folder : str 1282 The folder to save the network in. 1283 name : str 1284 The name to save the network under. 1285 1286 Returns 1287 ------- 1288 None 1289 1290 """ 1291 # If running a DDP only save with first thread 1292 if "RANK" in os.environ: 1293 if int(os.environ["RANK"]) != 0: 1294 return 1295 if not os.path.isdir(folder): 1296 os.makedirs(folder) 1297 save_point = folder + "/" 1298 if not os.path.isdir(save_point): 1299 os.mkdir(save_point) 1300 for param in net.parameters(): 1301 param.data = param.data.contiguous() 1302 if GPA.pc.get_using_safe_tensors(): 1303 if GPA.pc.get_weight_tying_experimental(): 1304 save_model_with_weight_tying(net, save_point + name + ".pt") 1305 else: 1306 # Strip the . so that the naming is the same for everywhere but it works with state_dict naming 1307 not_save = [ns.lstrip('.') for ns in GPA.pc.get_module_names_to_not_save()] 1308 state_dict = {k: v for k, v in net.state_dict().items() 1309 if not any(k.startswith(ns) for ns in not_save)} 1310 save_file(state_dict, save_point + name + ".pt") 1311 else: 1312 torch.save(net, save_point + name + ".pt")
Save the network
This is called within save_system after the tracker has been turned into a single tensor to be saved as a part of the network
Parameters
- net (nn.Module): The network to save.
- folder (str): The folder to save the network in.
- name (str): The name to save the network under.
Returns
- None
def
save_pai_net(net, folder, name):
1359def save_pai_net(net, folder, name): 1360 """Save the final pai network 1361 1362 This can be called after training to save the final network 1363 with all scaffolding removed so only the refined weights remain 1364 1365 Parameters 1366 ---------- 1367 net : nn.Module 1368 The network to save. 1369 folder : str 1370 The folder to save the network in. 1371 name : str 1372 The name to save the network under. 1373 1374 Returns 1375 ------- 1376 None 1377 1378 """ 1379 # if running a DDP only save with first thread 1380 if "RANK" in os.environ: 1381 if int(os.environ["RANK"]) != 0: 1382 return 1383 1384 # print('calling save: %s' % name) 1385 # GPA.pai_tracker.archive_layer() 1386 # These deep copys are required or the real model will also have its layers replaced 1387 net = prepare_final_model(net) 1388 if not os.path.isdir(folder): 1389 os.makedirs(folder) 1390 save_point = folder + "/" 1391 if not os.path.isdir(save_point): 1392 os.mkdir(save_point) 1393 1394 if GPA.pc.get_using_safe_tensors(): 1395 if GPA.pc.get_weight_tying_experimental(): 1396 save_model_with_weight_tying(net, save_point + name + "_pai.pt") 1397 else: 1398 save_file(net.state_dict(), save_point + name + "_pai.pt") 1399 else: 1400 torch.save(net, save_point + name + "_pai.pt")
Save the final pai network
This can be called after training to save the final network with all scaffolding removed so only the refined weights remain
Parameters
- net (nn.Module): The network to save.
- folder (str): The folder to save the network in.
- name (str): The name to save the network under.
Returns
- None
def
manual_load_state_dict(model, state_dict):
1403def manual_load_state_dict(model, state_dict): 1404 """Load a state dict into a model key-by-key with relaxed checks. 1405 1406 Parameters 1407 ---------- 1408 model : nn.Module 1409 Model receiving parameters. 1410 state_dict : dict 1411 Source parameter dictionary. 1412 1413 Notes 1414 ----- 1415 Keys configured in ``module_names_to_not_save`` are skipped. 1416 1417 Returns 1418 ------- 1419 None 1420 This function does not return a value. 1421 """ 1422 own_state = model.state_dict() 1423 not_save = [ns.lstrip('.') for ns in GPA.pc.get_module_names_to_not_save()] 1424 for name, param in state_dict.items(): 1425 if any(name.startswith(ns) for ns in not_save): 1426 print("skipping loading %s based on module_names_to_not_save" % name) 1427 continue 1428 if name not in own_state: 1429 print(f"Warning: {name} not found in model state_dict") 1430 continue 1431 if isinstance(param, torch.nn.Parameter): 1432 # Backwards compatibility for serialized parameters 1433 param = param.data 1434 try: 1435 own_state[name].copy_(param) 1436 except Exception as e: 1437 print(f"Error loading {name}: {e}") 1438 print("Manual load complete")
Load a state dict into a model key-by-key with relaxed checks.
Parameters
- model (nn.Module): Model receiving parameters.
- state_dict (dict): Source parameter dictionary.
Notes
Keys configured in module_names_to_not_save are skipped.
Returns
- None: This function does not return a value.
def
load_net(net, folder, name):
1441def load_net(net, folder, name): 1442 """load the network 1443 1444 This is called within load_system after the tracker has been 1445 loaded 1446 1447 Parameters 1448 ---------- 1449 net : nn.Module 1450 The network to save. 1451 folder : str 1452 The folder to save the network in. 1453 name : str 1454 The name to save the network under. 1455 1456 Returns 1457 ------- 1458 nn.Module 1459 The loaded network. 1460 1461 """ 1462 save_point = folder + "/" 1463 if GPA.pc.get_using_safe_tensors(): 1464 model_path = save_point + name + ".pt" 1465 if GPA.pc.get_weight_tying_experimental(): 1466 return load_model_with_weight_tying(net, model_path) 1467 else: 1468 try: 1469 with safe_open(model_path, framework="pt") as f: 1470 metadata = f.metadata() 1471 if metadata and "weight_tying" in metadata: 1472 return load_model_with_weight_tying(net, model_path) 1473 except Exception: 1474 pass 1475 state_dict = load_file(model_path) 1476 else: 1477 # Different versions of torch require this change 1478 try: 1479 state_dict = torch.load( 1480 save_point + name + ".pt", 1481 map_location=torch.device("cpu"), 1482 weights_only=False, 1483 ).state_dict() 1484 except: 1485 try: 1486 state_dict = torch.load( 1487 save_point + name + ".pt", map_location=torch.device("cpu") 1488 ).state_dict() 1489 except: 1490 state_dict = torch.load( 1491 save_point + name + ".pt", map_location=torch.device("cpu") 1492 ) 1493 return load_net_from_dict(net, state_dict)
load the network
This is called within load_system after the tracker has been loaded
Parameters
- net (nn.Module): The network to save.
- folder (str): The folder to save the network in.
- name (str): The name to save the network under.
Returns
- nn.Module: The loaded network.
def
get_module_base_name(module):
1496def get_module_base_name(module): 1497 """Normalize a wrapped module name for state-dict key lookup. 1498 1499 Parameters 1500 ---------- 1501 module : nn.Module 1502 Module containing a ``name`` attribute. 1503 1504 Returns 1505 ------- 1506 str 1507 Base name with leading dot and optional ``module.`` prefix removed. 1508 """ 1509 module_name = module.name 1510 # This should always be true 1511 if module_name[0] == ".": 1512 # strip "." 1513 module_name = module_name[1:] 1514 # If it was a dataparallel it will also have a module at the start 1515 # so strip that for loading 1516 if module_name[:6] == "module": 1517 module_name = module_name[7:] 1518 return module_name
Normalize a wrapped module name for state-dict key lookup.
Parameters
- module (nn.Module):
Module containing a
nameattribute.
Returns
- str: Base name with leading dot and optional
module.prefix removed.
def
load_net_from_dict(net, state_dict):
1521def load_net_from_dict(net, state_dict): 1522 """load the network 1523 1524 This is called within load_net 1525 1526 Parameters 1527 ---------- 1528 net : nn.Module 1529 The network to save. 1530 state_dict : dict 1531 The state dictionary to load. 1532 1533 Returns 1534 ------- 1535 nn.Module 1536 The loaded network. 1537 1538 """ 1539 if GPA.pc.get_verbose(): 1540 print("loading net from dict") 1541 pai_modules = get_pai_modules(net, 0) 1542 if pai_modules == []: 1543 print( 1544 "PAI load_net and load_system uses a state_dict so it must be\n" 1545 "called with a net after perforate_model has been called" 1546 ) 1547 print( 1548 "This is being flagged because you are attempting to load a model\n" 1549 "that does not have any pai_modules in it. Confirm that you are calling\n" 1550 "perforate_model on the correct model, and the same model is the one\n" 1551 "being passed into add_validation_score" 1552 ) 1553 import pdb # This needs to be here for cython for some reason. 1554 pdb.set_trace() 1555 sys.exit(-1) 1556 if GPA.pc.get_verbose(): 1557 print( 1558 "setting up arrays and simulating cycles for %d pai modules" 1559 % len(pai_modules) 1560 ) 1561 not_save = GPA.pc.get_module_names_to_not_save() 1562 for module in pai_modules: 1563 if any(module.name.startswith(ns) for ns in not_save): 1564 print("skipping loading %s based on module_names_to_not_save" % module.name) 1565 continue 1566 # Set up name to be what will be saved in the state dict 1567 module_name = get_module_base_name(module) 1568 module.clear_dendrites() 1569 for tracker in module.dendrite_module.dendrite_values: 1570 try: 1571 tracker.setup_arrays( 1572 len( 1573 state_dict[ 1574 module_name + ".dendrite_module.dendrite_values.0.shape" 1575 ] 1576 ) 1577 ) 1578 except Exception as e: 1579 print(e) 1580 print( 1581 "This value is missing from the state dict\n" 1582 "When missing this value it typically means you\n" 1583 "converted a module but didn't actually use it in\n" 1584 "your forward and backward pass." 1585 ) 1586 print("module was: %s" % module.name) 1587 print("There are many reasons this can happen:") 1588 print( 1589 "\n1 - check your model definition and forward function and " 1590 "ensure this module is being used properly" 1591 ) 1592 print( 1593 "with GPA.pc.set_verbose(True) you can confirm this is the case if\n" 1594 'you do not see a "setting d shape for" this module at the first training batch.' 1595 ) 1596 print( 1597 "If this is the case, and it is correct to not be passing data through it\n" 1598 "Set it to be a tracked module with:\n" 1599 'GPA.pc.append_module_ids_to_track(["%s"]) to leave it out ' 1600 % module.name 1601 ) 1602 print( 1603 "\n2 - This can happen if you adjusted your model " 1604 "definition after calling perforate_model" 1605 ) 1606 print( 1607 "for example with torch.compile. If the module name " 1608 "printed above does not contain all modules leading " 1609 "to the main definition" 1610 ) 1611 print( 1612 "this is likely the case for your problem. Fix by " 1613 "calling perforate_model after all other model " 1614 "initialization steps" 1615 ) 1616 first_key = next(iter(state_dict.keys())) 1617 print( 1618 "\n3 - This can happen is if the model where you called perforate_model\n" 1619 "and the model within add_validation_score are not the same. \n" 1620 "Check if the module above and .%s have the same prefix\n" 1621 % first_key 1622 ) 1623 print( 1624 "if one starts with .model or .base etc and the other does not, this is the problem." 1625 ) 1626 1627 print( 1628 "\n4 - If you are using this module but then not actually including\n" 1629 "the correct output tensor in the forward. For example\n" 1630 "if you are using an LSTM and forwarding hidden instead of otput\n" 1631 "but your processors are set up to work with output" 1632 ) 1633 print( 1634 "\n5 - if you are not properly calling backward at all." 1635 " If this is the first module in your network it is more" 1636 "likely this is the problem." 1637 "One check in these cases is to make sure you do not call an initial validation score" 1638 "before the first backward call.\nIf you do this, while testing_dendrite_capacity is True" 1639 "this error will be triggered." 1640 ) 1641 print( 1642 "\n6 - You have converted a module that is in a frozen" 1643 " part of the network and thus no gradients are flowing" 1644 ) 1645 print( 1646 "\n7 - You are running multiple experiments at once with the same save_name." 1647 " When running concurrent trials be sure to add save_name=<unique_name> to perforate_model." 1648 ) 1649 import pdb # This needs to be here for cython for some reason. 1650 pdb.set_trace() 1651 1652 # Perform as many cycles as the state dict has 1653 num_cycles = int(state_dict[module_name + ".dendrite_module.num_cycles"].item()) 1654 if num_cycles > 0: 1655 simulate_cycles(module, num_cycles, doing_pai=True) 1656 # Handle tracker_string loading with flexible key matching 1657 tracker_key = None 1658 if "tracker_string" in state_dict: 1659 tracker_key = "tracker_string" 1660 else: 1661 # Search for keys containing "tracker_string" 1662 tracker_keys = [key for key in state_dict.keys() if "tracker_string" in key] 1663 if len(tracker_keys) == 1: 1664 tracker_key = tracker_keys[0] 1665 elif len(tracker_keys) > 1: 1666 print(f"Error: Multiple tracker_string keys found: {tracker_keys}") 1667 import pdb # This needs to be here for cython for some reason. 1668 pdb.set_trace() 1669 else: 1670 print("Error: No tracker_string found in state_dict") 1671 import pdb # This needs to be here for cython for some reason. 1672 pdb.set_trace() 1673 1674 if hasattr(net, "tracker_string"): 1675 net.tracker_string = state_dict[tracker_key] 1676 else: 1677 net.register_buffer("tracker_string", state_dict[tracker_key]) 1678 try: 1679 load_result = net.load_state_dict(state_dict, strict=False) 1680 not_save_state_names = [ns.lstrip('.') for ns in not_save] 1681 1682 def is_ignored_key(key): 1683 """Check whether a state-dict key should be ignored. 1684 1685 Parameters 1686 ---------- 1687 key : str 1688 State-dict key to test. 1689 1690 Returns 1691 ------- 1692 bool 1693 ``True`` when key belongs to a not-saved namespace. 1694 """ 1695 return any(key.startswith(ns) for ns in not_save_state_names) 1696 1697 missing_keys = [key for key in load_result.missing_keys if not is_ignored_key(key)] 1698 unexpected_keys = [key for key in load_result.unexpected_keys if not is_ignored_key(key)] 1699 1700 if GPA.pc.get_strict_loading() and (missing_keys or unexpected_keys): 1701 raise RuntimeError( 1702 "Error(s) in loading state_dict for %s:\n\tMissing key(s) in state_dict: %s. \n\tUnexpected key(s) in state_dict: %s." 1703 % (type(net).__name__, missing_keys, unexpected_keys) 1704 ) 1705 except Exception as e: 1706 """ 1707 When modules have high depth to them (i.e. modules within modules not number of layers) 1708 PyTorch can have trouble loading state dicts even when they are correct. 1709 This is a workaround to manually load the state dict if this happens. 1710 """ 1711 filtered_net_keys = { 1712 key 1713 for key in net.state_dict().keys() 1714 if not any(key.startswith(ns.lstrip('.')) for ns in not_save) 1715 } 1716 if filtered_net_keys == set(state_dict.keys()): 1717 print("Attempting manual loading of state_dict") 1718 manual_load_state_dict(net, state_dict) 1719 else: 1720 print(f"Error loading state_dict: {e}") 1721 print("If the error is due to missing keys (e.g., from code changes), you can try:") 1722 print(" GPA.pc.set_strict_loading(False)") 1723 print(" Do not change this unless you are certain the missing keys are not important to load and are expected due to code changes or arch changes.") 1724 print("\ntype 'c' to print full state dicts\n") 1725 import pdb # This needs to be here for cython for some reason. 1726 pdb.set_trace() 1727 print("net state dict is:") 1728 print(net.state_dict()) 1729 print("loaded state dict is:") 1730 print(state_dict) 1731 print( 1732 "Try to check differences. Likely is caused by a module not " 1733 "being converted that should be or vice versa" 1734 ) 1735 pdb.set_trace() 1736 net.to(GPA.pc.get_device()) 1737 return net
load the network
This is called within load_net
Parameters
- net (nn.Module): The network to save.
- state_dict (dict): The state dictionary to load.
Returns
- nn.Module: The loaded network.
def
pai_save_system(net, folder, name):
1740def pai_save_system(net, folder, name): 1741 """Save the entire system with scaffolding removed 1742 1743 This is used for the final network for inference after training 1744 1745 Parameters 1746 ---------- 1747 net : nn.Module 1748 The network to save. 1749 folder : str 1750 The folder to save the network in. 1751 name : str 1752 The name to save the network under. 1753 1754 Returns 1755 ------- 1756 None 1757 1758 """ 1759 net.member_vars = {} 1760 for member_var in GPA.pai_tracker.member_vars: 1761 if member_var == "scheduler_instance" or member_var == "optimizer_instance": 1762 continue 1763 net.member_vars[member_var] = GPA.pai_tracker.member_vars[member_var] 1764 pai_save_net(net, folder, name)
Save the entire system with scaffolding removed
This is used for the final network for inference after training
Parameters
- net (nn.Module): The network to save.
- folder (str): The folder to save the network in.
- name (str): The name to save the network under.
Returns
- None
def
deep_copy_pai(net):
1767def deep_copy_pai(net): 1768 """Deep copy a PAI network 1769 1770 1771 Parameters 1772 ---------- 1773 net : nn.Module 1774 The network to copy. 1775 1776 Returns 1777 ------- 1778 nn.Module 1779 The copied network. 1780 1781 Notes 1782 ---- 1783 This is required because processors must be cleared before calling copy 1784 1785 """ 1786 # Dont check this stuff if its before the perforate_model has been called and you're just copying a regular model 1787 if(GPA.pai_tracker != []): 1788 # Clear gradients before saving the model 1789 if ((GPA.pai_tracker.member_vars["optimizer_instance"]) is not None) and ( 1790 GPA.pai_tracker.member_vars["optimizer_instance"] != [] 1791 ): 1792 GPA.pai_tracker.member_vars["optimizer_instance"].zero_grad() 1793 GPA.pai_tracker.clear_all_processors() 1794 return copy.deepcopy(net)
Deep copy a PAI network
Parameters
- net (nn.Module): The network to copy.
Returns
- nn.Module: The copied network.
Notes
This is required because processors must be cleared before calling copy
def
prepare_final_model(net):
1797def prepare_final_model(net): 1798 """Prepare model for final save by removing scaffolding. 1799 1800 This performs all cleanup steps to convert a PAI model with scaffolding 1801 into a clean final model ready for inference or distribution. 1802 1803 Parameters 1804 ---------- 1805 net : nn.Module 1806 The network to prepare. 1807 1808 Returns 1809 ------- 1810 nn.Module 1811 The cleaned model with scaffolding removed. 1812 """ 1813 # Deep copy and clean the model (removes scaffolding) 1814 net = deep_copy_pai(net) 1815 net = BPA.blockwise_network(net) 1816 net = deep_copy_pai(net) 1817 net = CL.refresh_net(net) 1818 1819 # Remove tracker_string (not needed for final model) 1820 if hasattr(net, "tracker_string"): 1821 del net.tracker_string 1822 1823 # Make parameters contiguous 1824 for param in net.parameters(): 1825 param.data = param.data.contiguous() 1826 1827 return net
Prepare model for final save by removing scaffolding.
This performs all cleanup steps to convert a PAI model with scaffolding into a clean final model ready for inference or distribution.
Parameters
- net (nn.Module): The network to prepare.
Returns
- nn.Module: The cleaned model with scaffolding removed.
def
pai_save_net(net, folder, name):
1830def pai_save_net(net, folder, name): 1831 """Save the entire system with scaffolding removed 1832 1833 This is called within pai_save_system after the tracker has been 1834 turned into a single tensor to be saved as a part of the network 1835 1836 1837 Parameters 1838 ---------- 1839 net : nn.Module 1840 The network to save. 1841 folder : str 1842 The folder to save the network in. 1843 name : str 1844 The name to save the network under. 1845 1846 Returns 1847 ------- 1848 None 1849 1850 Notes 1851 ---- 1852 For open source implementation this is not as important since 1853 minimal values are already being used. 1854 1855 """ 1856 1857 if GPA.pc.get_perforated_backpropagation(): 1858 UPB.pb_save_net(net, folder, name) 1859 else: 1860 return
Save the entire system with scaffolding removed
This is called within pai_save_system after the tracker has been turned into a single tensor to be saved as a part of the network
Parameters
- net (nn.Module): The network to save.
- folder (str): The folder to save the network in.
- name (str): The name to save the network under.
Returns
- None
Notes
For open source implementation this is not as important since minimal values are already being used.
def
simulate_cycles(module, num_cycles, doing_pai):
1863def simulate_cycles(module, num_cycles, doing_pai): 1864 """Simulate dendrite addition cycles 1865 1866 Simulate the back and forth processes of adding dendrites to build a 1867 pretrained dendrite model before loading weights. Required for loading 1868 dendrite save files from non dendrite initial models. 1869 1870 Parameters 1871 ---------- 1872 module : PA.PAINeuronModule 1873 The module to simulate cycles on. 1874 num_cycles : int 1875 The number of cycles to simulate. 1876 doing_pai : bool 1877 Whether to actually do the simulation. 1878 1879 Returns 1880 ------- 1881 None 1882 1883 """ 1884 1885 check_skipped = GPA.pc.get_checked_skipped_modules() 1886 if doing_pai is False: 1887 return 1888 GPA.pc.set_checked_skipped_modules(True) 1889 mode = "n" 1890 for i in range(num_cycles): 1891 if mode == "n": 1892 module.set_mode("p") 1893 module.create_new_dendrite_module() 1894 mode = "p" 1895 else: 1896 module.set_mode("n") 1897 mode = "n" 1898 GPA.pc.set_checked_skipped_modules(check_skipped)
Simulate dendrite addition cycles
Simulate the back and forth processes of adding dendrites to build a pretrained dendrite model before loading weights. Required for loading dendrite save files from non dendrite initial models.
Parameters
- module (PA.PAINeuronModule): The module to simulate cycles on.
- num_cycles (int): The number of cycles to simulate.
- doing_pai (bool): Whether to actually do the simulation.
Returns
- None
def
count_params(net):
1901def count_params(net): 1902 """Count the number of parameters in the network 1903 1904 If doing perforated backpropagation this calls the PB function 1905 which does not count scaffolding parameters since the final model 1906 will not have them. 1907 1908 Parameters 1909 ---------- 1910 net : nn.Module 1911 The network to count parameters in. 1912 1913 Returns 1914 ------- 1915 int 1916 The number of parameters in the network. 1917 1918 """ 1919 if GPA.pc.get_perforated_backpropagation(): 1920 return UPB.pb_count_params(net) 1921 parameters = net.named_parameters() 1922 unique_params = { 1923 p.data_ptr(): p for name, p in parameters if "parent_module" not in name 1924 }.values() 1925 return sum(p.numel() for p in unique_params)
Count the number of parameters in the network
If doing perforated backpropagation this calls the PB function which does not count scaffolding parameters since the final model will not have them.
Parameters
- net (nn.Module): The network to count parameters in.
Returns
- int: The number of parameters in the network.
def
change_learning_modes(net, folder, name, doing_pai):
1928def change_learning_modes(net, folder, name, doing_pai): 1929 """Change between neuron and dendrite learning modes 1930 1931 High level steps for entire system to switch back and forth between 1932 neuron learning and dendrite learning 1933 1934 Parameters 1935 ---------- 1936 net : nn.Module 1937 The network to change modes on. 1938 folder : str 1939 The folder to save/load the network in/from. 1940 name : str 1941 The name to save/load the network under. 1942 doing_pai : bool 1943 Whether to add dendrites when changing modes. 1944 1945 Returns 1946 ------- 1947 int 1948 The number of parameters in the network. 1949 1950 Notes 1951 ----- 1952 If doing_pai is False this just allows training to continue longer rather than early stopping 1953 1954 """ 1955 # If not adding dendrites this just allows training to continue longer with flags 1956 # every time early stopping should be occurring 1957 if doing_pai is False: 1958 GPA.pai_tracker.member_vars["switch_epochs"].append( 1959 GPA.pai_tracker.member_vars["num_epochs_run"] 1960 ) 1961 GPA.pai_tracker.member_vars["last_switch"] = GPA.pai_tracker.member_vars[ 1962 "switch_epochs" 1963 ][-1] 1964 GPA.pai_tracker.reset_vals_for_score_reset() 1965 return net 1966 if GPA.pai_tracker.member_vars["mode"] == "n": 1967 current_epoch = GPA.pai_tracker.member_vars["num_epochs_run"] 1968 overwritten_epochs = GPA.pai_tracker.member_vars["overwritten_epochs"] 1969 overwritten_extra = GPA.pai_tracker.member_vars["extra_scores"] 1970 if GPA.pc.get_drawing_pai(): 1971 overwritten_val = GPA.pai_tracker.member_vars["accuracies"] 1972 else: 1973 overwritten_val = GPA.pai_tracker.member_vars["neuron_accuracies"] 1974 """ 1975 If true don't load the best system 1976 because it will delete dendrites if the previous best was better than 1977 the current best 1978 """ 1979 if not GPA.pc.get_silent(): 1980 print("Importing best Model for switch to PA...") 1981 net = load_system(net, folder, name, switch_call=True) 1982 GPA.pai_tracker.set_dendrite_training() 1983 GPA.pai_tracker.member_vars["overwritten_epochs"] = overwritten_epochs 1984 GPA.pai_tracker.member_vars["overwritten_epochs"] += ( 1985 current_epoch - GPA.pai_tracker.member_vars["num_epochs_run"] 1986 ) 1987 GPA.pai_tracker.member_vars["total_epochs_run"] = ( 1988 GPA.pai_tracker.member_vars["num_epochs_run"] 1989 + GPA.pai_tracker.member_vars["overwritten_epochs"] 1990 ) 1991 1992 if GPA.pc.get_save_old_graph_scores(): 1993 GPA.pai_tracker.member_vars["overwritten_extras"].append(overwritten_extra) 1994 GPA.pai_tracker.member_vars["overwritten_vals"].append(overwritten_val) 1995 else: 1996 GPA.pai_tracker.member_vars["overwritten_extras"] = [overwritten_extra] 1997 GPA.pai_tracker.member_vars["overwritten_vals"] = [overwritten_val] 1998 if GPA.pc.get_drawing_pai(): 1999 GPA.pai_tracker.member_vars["n_switch_epochs"].append( 2000 GPA.pai_tracker.member_vars["num_epochs_run"] 2001 ) 2002 else: 2003 if len(GPA.pai_tracker.member_vars["switch_epochs"]) == 0: 2004 GPA.pai_tracker.member_vars["n_switch_epochs"].append( 2005 GPA.pai_tracker.member_vars["num_epochs_run"] 2006 ) 2007 else: 2008 GPA.pai_tracker.member_vars["n_switch_epochs"].append( 2009 GPA.pai_tracker.member_vars["n_switch_epochs"][-1] 2010 + ( 2011 (GPA.pai_tracker.member_vars["num_epochs_run"]) 2012 - (GPA.pai_tracker.member_vars["switch_epochs"][-1]) 2013 ) 2014 ) 2015 2016 GPA.pai_tracker.member_vars["switch_epochs"].append( 2017 GPA.pai_tracker.member_vars["num_epochs_run"] 2018 ) 2019 GPA.pai_tracker.member_vars["last_switch"] = GPA.pai_tracker.member_vars[ 2020 "switch_epochs" 2021 ][-1] 2022 2023 # Because open source version is only doing neuron training for 2024 # gradient descent dendrites, switch back to n mode right away 2025 if ( 2026 not GPA.pc.get_perforated_backpropagation() 2027 ) or GPA.pc.get_no_extra_n_modes(): 2028 net = change_learning_modes(net, folder, name, doing_pai) 2029 else: 2030 if not GPA.pc.get_silent(): 2031 print("Switching back to N...") 2032 set_best = GPA.pai_tracker.member_vars["current_n_set_global_best"] 2033 GPA.pai_tracker.set_neuron_training() 2034 if len(GPA.pai_tracker.member_vars["p_switch_epochs"]) == 0: 2035 GPA.pai_tracker.member_vars["p_switch_epochs"].append( 2036 ( 2037 (GPA.pai_tracker.member_vars["num_epochs_run"] - 1) 2038 - (GPA.pai_tracker.member_vars["switch_epochs"][-1]) 2039 ) 2040 ) 2041 else: 2042 GPA.pai_tracker.member_vars["p_switch_epochs"].append( 2043 GPA.pai_tracker.member_vars["p_switch_epochs"][-1] 2044 + ( 2045 (GPA.pai_tracker.member_vars["num_epochs_run"]) 2046 - (GPA.pai_tracker.member_vars["switch_epochs"][-1]) 2047 ) 2048 ) 2049 GPA.pai_tracker.member_vars["switch_epochs"].append( 2050 GPA.pai_tracker.member_vars["num_epochs_run"] 2051 ) 2052 GPA.pai_tracker.member_vars["last_switch"] = GPA.pai_tracker.member_vars[ 2053 "switch_epochs" 2054 ][-1] 2055 # Will be false for open source implementation 2056 if GPA.pc.get_retain_all_dendrites() or ( 2057 GPA.pc.get_learn_dendrites_live() and set_best 2058 ): 2059 if not GPA.pc.get_silent(): 2060 print( 2061 "Saving model before starting normal training to " 2062 "retain PBNodes regardless of next N Phase results" 2063 ) 2064 save_system(net, folder, name) 2065 # if its just doing P for learn PAI live then switch back immediately 2066 if GPA.pc.get_perforated_backpropagation() and GPA.pc.get_no_extra_n_modes(): 2067 net = change_learning_modes(net, folder, name, doing_pai) 2068 2069 GPA.pai_tracker.member_vars["param_counts"].append(count_params(net)) 2070 2071 return net
Change between neuron and dendrite learning modes
High level steps for entire system to switch back and forth between neuron learning and dendrite learning
Parameters
- net (nn.Module): The network to change modes on.
- folder (str): The folder to save/load the network in/from.
- name (str): The name to save/load the network under.
- doing_pai (bool): Whether to add dendrites when changing modes.
Returns
- int: The number of parameters in the network.
Notes
If doing_pai is False this just allows training to continue longer rather than early stopping
def
find_param_name_by_id(model, param_id):
2074def find_param_name_by_id(model, param_id): 2075 """ 2076 This is only used for debugging. 2077 Return the fully-qualified parameter name (e.g. "layer1.conv.weight") 2078 for the parameter whose id matches param_id. Returns None if not found. 2079 2080 This uses model.named_parameters(), which already recurses through submodules. 2081 2082 Parameters 2083 ---------- 2084 model : Model to look for param id. 2085 param_id : pointer to a parameter. 2086 2087 Returns 2088 ------- 2089 String representing the module within the model or None if not found. 2090 """ 2091 for name, p in model.named_parameters(recurse=True): 2092 if id(p) == param_id: 2093 return "." + name 2094 return None
This is only used for debugging. Return the fully-qualified parameter name (e.g. "layer1.conv.weight") for the parameter whose id matches param_id. Returns None if not found.
This uses model.named_parameters(), which already recurses through submodules.
Parameters
model (Model to look for param id.):
param_id (pointer to a parameter.):
Returns
- String representing the module within the model or None if not found.
def
add_method_delegation_to_module(wrapper_module, method_name):
2097def add_method_delegation_to_module(wrapper_module, method_name): 2098 """Add delegating methods to a wrapper module that has a main_module attribute. 2099 2100 This adds the specified methods to the wrapper module instance so they 2101 properly delegate to the wrapped main_module. Works for any wrapper module 2102 (TrackedNeuronModule, PAINeuronModule, etc.) that has a main_module attribute. 2103 2104 Args: 2105 wrapper_module: A wrapper module instance with a main_module attribute 2106 method_name: The method name to delegate (e.g., '_gradient_checkpointing_func') 2107 2108 Parameters 2109 ---------- 2110 wrapper_module : PyTorch Module that contains a sub module. 2111 method_name : Name of method to wrap. 2112 2113 Returns 2114 ------- 2115 None 2116 """ 2117 import types 2118 2119 if hasattr(wrapper_module.main_module, method_name): 2120 # Create a delegating method that forwards to main_module 2121 def make_delegated_method(name): 2122 """Create a bound delegation function for a given attribute name. 2123 2124 Parameters 2125 ---------- 2126 name : str 2127 Attribute name to forward to ``main_module``. 2128 2129 Returns 2130 ------- 2131 callable 2132 Function that delegates access or invocation. 2133 """ 2134 def delegated_method(self, *args, **kwargs): 2135 """Delegate attribute access or method call to ``main_module``. 2136 2137 Parameters 2138 ---------- 2139 *args : tuple 2140 Positional arguments forwarded to delegated callables. 2141 **kwargs : dict 2142 Keyword arguments forwarded to delegated callables. 2143 2144 Returns 2145 ------- 2146 Any 2147 Delegated attribute value or method result. 2148 """ 2149 main_module_attr = getattr(self.main_module, name, None) 2150 if main_module_attr is None: 2151 raise AttributeError( 2152 f"'{type(self.main_module).__name__}' object has no attribute '{name}'" 2153 ) 2154 if callable(main_module_attr): 2155 return main_module_attr(*args, **kwargs) 2156 return main_module_attr 2157 2158 return delegated_method 2159 2160 # Bind it to this specific instance 2161 setattr( 2162 wrapper_module, 2163 method_name, 2164 types.MethodType(make_delegated_method(method_name), wrapper_module), 2165 )
Add delegating methods to a wrapper module that has a main_module attribute.
This adds the specified methods to the wrapper module instance so they properly delegate to the wrapped main_module. Works for any wrapper module (TrackedNeuronModule, PAINeuronModule, etc.) that has a main_module attribute.
Args: wrapper_module: A wrapper module instance with a main_module attribute method_name: The method name to delegate (e.g., '_gradient_checkpointing_func')
Parameters
wrapper_module (PyTorch Module that contains a sub module.):
method_name (Name of method to wrap.):
Returns
- None
def
apply_method_delegation_to_model(model, method_name, main_module_type):
2168def apply_method_delegation_to_model(model, method_name, main_module_type): 2169 """Recursively apply method delegation to all wrapper modules with main_module in a model. 2170 2171 This traverses the entire model and adds method delegation for any module that has 2172 a main_module attribute and optionally matches specified types. 2173 2174 Args: 2175 model: The PyTorch model to traverse 2176 method_name: The method name to delegate (e.g., '_gradient_checkpointing_func') 2177 main_module_type: main_module type name to filter by. 2178 Example: 'Qwen2DecoderLayer' 2179 2180 Example: 2181 # Apply gradient checkpointing delegation to all decoder layers 2182 apply_method_delegation_to_model( 2183 model, 2184 '_gradient_checkpointing_func', 2185 main_module_type='Qwen2DecoderLayer' 2186 ) 2187 2188 Parameters 2189 ---------- 2190 model : PyTorch model. 2191 method_name : method to delegate. 2192 main_module_type : type of module that has this method. 2193 2194 Returns 2195 ------- 2196 None 2197 This function does not return a value. 2198 """ 2199 count = 0 2200 for name, module in model.named_modules(): 2201 # Check if module has main_module attribute (it's a wrapper) 2202 if hasattr(module, "main_module"): 2203 # Check if we should apply based on main_module type 2204 should_apply = True 2205 if main_module_type is not None: 2206 main_module_type_name = type(module.main_module).__name__ 2207 should_apply = main_module_type_name == main_module_type 2208 2209 if should_apply: 2210 add_method_delegation_to_module(module, method_name) 2211 count += 1 2212 2213 print(f"[PAI] Applied method delegation to {count} wrapper module instances")
Recursively apply method delegation to all wrapper modules with main_module in a model.
This traverses the entire model and adds method delegation for any module that has a main_module attribute and optionally matches specified types.
Args: model: The PyTorch model to traverse method_name: The method name to delegate (e.g., '_gradient_checkpointing_func') main_module_type: main_module type name to filter by. Example: 'Qwen2DecoderLayer'
Example: # Apply gradient checkpointing delegation to all decoder layers apply_method_delegation_to_model( model, '_gradient_checkpointing_func', main_module_type='Qwen2DecoderLayer' )
Parameters
model (PyTorch model.):
method_name (method to delegate.):
main_module_type (type of module that has this method.):
Returns
- None: This function does not return a value.
def
make_json_serializable(obj):
2216def make_json_serializable(obj): 2217 """Recursively convert non-JSON-serializable objects to strings. 2218 2219 Parameters 2220 ---------- 2221 obj : any 2222 The object to convert 2223 2224 Returns 2225 ------- 2226 Any 2227 JSON-serializable version of the object 2228 """ 2229 if isinstance(obj, (str, int, float, bool, type(None))): 2230 return obj 2231 elif isinstance(obj, dict): 2232 return {k: make_json_serializable(v) for k, v in obj.items()} 2233 elif isinstance(obj, (list, tuple)): 2234 return [make_json_serializable(item) for item in obj] 2235 else: 2236 # Convert non-serializable types to string 2237 return str(obj)
Recursively convert non-JSON-serializable objects to strings.
Parameters
- obj (any): The object to convert
Returns
- Any: JSON-serializable version of the object
def
extract_gpa_config():
2240def extract_gpa_config(): 2241 """Extract all configuration from GPA.pc by calling all get_* methods. 2242 2243 Returns 2244 ------- 2245 dict[str, Any] 2246 Dictionary with all GPA.pc configuration values and type metadata 2247 2248 Examples 2249 -------- 2250 >>> config = extract_gpa_config() 2251 >>> # Returns: {'max_dendrites': 10, 'device': 'cuda', '_types': {...}} 2252 2253 Parameters 2254 ---------- 2255 None 2256 2257 """ 2258 config = {} 2259 config_types = {} 2260 2261 # Get all attributes from GPA.pc 2262 for attr_name in dir(GPA.pc): 2263 # Check if it starts with 'get_' 2264 if attr_name.startswith("get_"): 2265 try: 2266 # Get the method 2267 method = getattr(GPA.pc, attr_name) 2268 2269 # Check if it's callable 2270 if callable(method): 2271 # Call it and store result with key as name without 'get_' 2272 key = attr_name[4:] # Remove 'get_' prefix 2273 value = method() 2274 2275 # Check if this is an array (has corresponding append_ method) 2276 append_method_name = f"append_{key}" 2277 is_array = hasattr(GPA.pc, append_method_name) 2278 2279 if is_array and isinstance(value, (list, tuple)): 2280 # Store array element type 2281 if len(value) > 0: 2282 element_type = type(value[0]).__name__ 2283 else: 2284 element_type = None # empty array, no conversion needed 2285 config_types[key] = { 2286 "is_array": True, 2287 "element_type": element_type, 2288 } 2289 else: 2290 # Store value type 2291 config_types[key] = { 2292 "is_array": False, 2293 "type": type(value).__name__, 2294 } 2295 2296 # Make sure value is JSON serializable 2297 config[key] = make_json_serializable(value) 2298 except Exception as e: 2299 # Skip if method fails 2300 if GPA.pc.get_verbose(): 2301 print(f"Skipping {attr_name}: {e}") 2302 continue 2303 2304 # Add types metadata to config 2305 config["_types"] = config_types 2306 2307 return config
Extract all configuration from GPA.pc by calling all get_* methods.
Returns
- dict[str, Any]: Dictionary with all GPA.pc configuration values and type metadata
Examples
>>> config = extract_gpa_config()
>>> # Returns: {'max_dendrites': 10, 'device': 'cuda', '_types': {...}}
Parameters
- None
def
convert_to_type(value, type_name):
2310def convert_to_type(value, type_name): 2311 """Convert a value to the specified type. 2312 2313 Parameters 2314 ---------- 2315 value : any 2316 The value to convert 2317 type_name : str 2318 The target type name 2319 2320 Returns 2321 ------- 2322 Any 2323 The converted value 2324 """ 2325 if type_name == "NoneType" or value is None: 2326 return None 2327 elif type_name == "bool": 2328 if isinstance(value, str): 2329 return value.lower() in ("true", "1", "yes") 2330 return bool(value) 2331 elif type_name == "int": 2332 return int(value) 2333 elif type_name == "float": 2334 return float(value) 2335 elif type_name == "str": 2336 return str(value) 2337 elif type_name == "list": 2338 if not isinstance(value, list): 2339 return [value] 2340 return value 2341 elif type_name == "dict": 2342 if not isinstance(value, dict): 2343 return {} 2344 return value 2345 elif type_name == "type": 2346 # Handle type objects - convert string representation back to type 2347 if isinstance(value, str): 2348 # Try to evaluate the type string (e.g., "<class 'torch.nn.Linear'>") 2349 # Extract the class path from the string 2350 if value.startswith("<class '") and value.endswith("'>"): 2351 class_path = value[ 2352 8:-2 2353 ] # Extract 'torch.nn.Linear' from "<class 'torch.nn.Linear'>" 2354 parts = class_path.split(".") 2355 # Try to import and get the type 2356 try: 2357 module_name = ".".join(parts[:-1]) 2358 class_name = parts[-1] 2359 module = __import__(module_name, fromlist=[class_name]) 2360 return getattr(module, class_name) 2361 except Exception as e: 2362 print( 2363 f"Warning: Could not convert type string '{value}' to actual type: {e}" 2364 ) 2365 return value 2366 return value 2367 return value 2368 elif type_name == "dtype": 2369 # Handle torch dtype objects 2370 if isinstance(value, str): 2371 # Convert string like "torch.float32" to actual dtype 2372 import torch 2373 2374 try: 2375 # Try to get the dtype from torch module 2376 if value.startswith("torch."): 2377 dtype_name = value.split(".")[ 2378 1 2379 ] # Get 'float32' from 'torch.float32' 2380 return getattr(torch, dtype_name) 2381 else: 2382 return getattr(torch, value) 2383 except Exception as e: 2384 print( 2385 f"Warning: Could not convert dtype string '{value}' to actual dtype: {e}" 2386 ) 2387 return value 2388 return value 2389 elif type_name == "device": 2390 # Handle torch device objects 2391 if isinstance(value, str): 2392 # Convert string like "cuda" or "cpu" to torch.device 2393 import torch 2394 2395 try: 2396 return torch.device(value) 2397 except Exception as e: 2398 print( 2399 f"Warning: Could not convert device string '{value}' to actual device: {e}" 2400 ) 2401 return value 2402 return value 2403 elif type_name == "builtin_function_or_method": 2404 # Handle torch functions like torch.sigmoid, torch.relu, etc. 2405 if isinstance(value, str): 2406 # Parse string like "<built-in method sigmoid of type object at 0x...>" 2407 # to extract the function name 2408 import torch 2409 2410 try: 2411 if "<built-in method " in value and " of type object" in value: 2412 # Extract function name between '<built-in method ' and ' of type object' 2413 start = value.find("<built-in method ") + len("<built-in method ") 2414 end = value.find(" of type object") 2415 func_name = value[start:end] 2416 # Try to get the function from torch module 2417 if hasattr(torch, func_name): 2418 return getattr(torch, func_name) 2419 else: 2420 print(f"Warning: torch.{func_name} not found") 2421 return value 2422 else: 2423 return value 2424 except Exception as e: 2425 print( 2426 f"Warning: Could not convert builtin function string '{value}': {e}" 2427 ) 2428 return value 2429 return value 2430 else: 2431 # Unknown type - error and debug 2432 print(f"ERROR: Unknown type '{type_name}' for value: {value}") 2433 print(f"Type of value is: {type(value).__name__}") 2434 pdb.set_trace() 2435 return value
Convert a value to the specified type.
Parameters
- value (any): The value to convert
- type_name (str): The target type name
Returns
- Any: The converted value
def
convert_to_type_array(value, element_type):
2438def convert_to_type_array(value, element_type): 2439 """Convert an array's elements to the specified type. 2440 2441 Parameters 2442 ---------- 2443 value : list or tuple 2444 The array to convert 2445 element_type : str or None 2446 The target type name for elements, None if array was empty 2447 2448 Returns 2449 ------- 2450 list 2451 The array with converted elements 2452 """ 2453 if not isinstance(value, (list, tuple)): 2454 return value 2455 # If element_type is None (empty array), no conversion needed 2456 if element_type is None: 2457 return list(value) if isinstance(value, tuple) else value 2458 return [convert_to_type(item, element_type) for item in value]
Convert an array's elements to the specified type.
Parameters
- value (list or tuple): The array to convert
- element_type (str or None): The target type name for elements, None if array was empty
Returns
- list: The array with converted elements
def
set_gpa_config(config):
2461def set_gpa_config(config): 2462 """Set GPA.pc configuration by calling all set_* methods. 2463 2464 This is the reverse of extract_gpa_config(). It takes a configuration 2465 dictionary and calls the corresponding set_* methods on GPA.pc. 2466 Uses type metadata to ensure values are converted to the correct type. 2467 2468 Parameters 2469 ---------- 2470 config : dict 2471 Dictionary with configuration values (keys without 'set_' prefix) 2472 and optional '_types' metadata 2473 2474 Examples 2475 -------- 2476 >>> config = {'verbose': True, 'device': 'cuda'} 2477 >>> set_gpa_config(config) 2478 # Calls GPA.pc.set_verbose(True), GPA.pc.set_device('cuda'), etc. 2479 2480 Returns 2481 ------- 2482 Count of parameters that were set 2483 """ 2484 set_count = 0 2485 skip_count = 0 2486 2487 # Extract type information 2488 config_types = config.get("_types", {}) 2489 2490 for key, value in config.items(): 2491 # Skip the types metadata 2492 if key == "_types": 2493 continue 2494 2495 # Construct the set method name 2496 set_method_name = f"set_{key}" 2497 2498 # Check if the set method exists 2499 if hasattr(GPA.pc, set_method_name): 2500 try: 2501 method = getattr(GPA.pc, set_method_name) 2502 if callable(method): 2503 # Convert value to correct type if we have type info 2504 if key in config_types: 2505 type_info = config_types[key] 2506 if type_info.get("is_array", False): 2507 # Convert array elements to correct type 2508 element_type = type_info.get("element_type", "str") 2509 value = convert_to_type_array(value, element_type) 2510 else: 2511 # Convert single value to correct type 2512 value_type = type_info.get("type", "str") 2513 value = convert_to_type(value, value_type) 2514 2515 method(value) 2516 set_count += 1 2517 if GPA.pc.get_verbose(): 2518 print(f"Set {key} = {value}") 2519 except Exception as e: 2520 skip_count += 1 2521 if GPA.pc.get_verbose(): 2522 print(f"Failed to set {key}: {e}") 2523 else: 2524 skip_count += 1 2525 if GPA.pc.get_verbose(): 2526 print(f"No setter found for {key} (looking for {set_method_name})") 2527 2528 if GPA.pc.get_verbose(): 2529 print(f"Applied {set_count} PAI configuration settings ({skip_count} skipped)") 2530 2531 return set_count
Set GPA.pc configuration by calling all set_* methods.
This is the reverse of extract_gpa_config(). It takes a configuration dictionary and calls the corresponding set_* methods on GPA.pc. Uses type metadata to ensure values are converted to the correct type.
Parameters
- config (dict): Dictionary with configuration values (keys without 'set_' prefix) and optional '_types' metadata
Examples
>>> config = {'verbose': True, 'device': 'cuda'}
>>> set_gpa_config(config)
<h1 id="calls-gpapcset_verbosetrue-gpapcset_devicecuda-etc">Calls GPA.pc.set_verbose(True), GPA.pc.set_device('cuda'), etc.</h1>
Returns
- Count of parameters that were set
def
upload_to_huggingface(*args, **kwargs):
2742 def upload_to_huggingface(*args, **kwargs): 2743 """Raise an informative error when HuggingFace dependencies are missing. 2744 2745 Parameters 2746 ---------- 2747 *args : tuple[Any, ...] 2748 Positional arguments accepted for API compatibility. 2749 **kwargs : dict[str, Any] 2750 Keyword arguments accepted for API compatibility. 2751 2752 Returns 2753 ------- 2754 None 2755 Always raises ``ImportError``. 2756 """ 2757 raise ImportError( 2758 "huggingface_hub is required for upload_to_huggingface. " 2759 "Install it with: pip install huggingface_hub" 2760 )
Raise an informative error when HuggingFace dependencies are missing.
Parameters
- *args (tuple[Any, ...]): Positional arguments accepted for API compatibility.
- **kwargs (dict[str, Any]): Keyword arguments accepted for API compatibility.
Returns
- None: Always raises
ImportError.
def
from_hf_pretrained(*args, **kwargs):
2762 def from_hf_pretrained(*args, **kwargs): 2763 """Raise an informative error when HuggingFace dependencies are missing. 2764 2765 Parameters 2766 ---------- 2767 *args : tuple[Any, ...] 2768 Positional arguments accepted for API compatibility. 2769 **kwargs : dict[str, Any] 2770 Keyword arguments accepted for API compatibility. 2771 2772 Returns 2773 ------- 2774 None 2775 Always raises ``ImportError``. 2776 """ 2777 raise ImportError( 2778 "huggingface_hub is required for from_hf_pretrained. " 2779 "Install it with: pip install huggingface_hub" 2780 )
Raise an informative error when HuggingFace dependencies are missing.
Parameters
- *args (tuple[Any, ...]): Positional arguments accepted for API compatibility.
- **kwargs (dict[str, Any]): Keyword arguments accepted for API compatibility.
Returns
- None: Always raises
ImportError.