| 651 | self._replace_module(child, full_name, "") |
| 652 | |
| 653 | def _replace_module(self, r_module, prev_name='', prev_class_name=''): |
| 654 | if prev_name == '' and prev_class_name == '': |
| 655 | self._configure_gathered_column_tie_fallbacks() |
| 656 | |
| 657 | for name, child in r_module.named_children(): |
| 658 | if getattr(child, "_is_autoep_layer", False): |
| 659 | full_name = prev_name + '.' + name if prev_name else name |
| 660 | self._replace_autoep_shared_experts(child, full_name) |
| 661 | continue |
| 662 | |
| 663 | if prev_class_name == "": |
| 664 | class_name = prev_name |
| 665 | elif prev_name == "": |
| 666 | class_name = prev_class_name |
| 667 | else: |
| 668 | class_name = prev_class_name + '.' + prev_name |
| 669 | checking_key = self.prefix + '.' + class_name + '.' + name + '.' if class_name != "" else self.prefix + '.' + name + '.' |
| 670 | if Loading.is_load_module(child) and self.state_dict is not None: |
| 671 | if any(checking_key in item for item in self.state_dict): |
| 672 | Loading.load(child, self.state_dict, checking_key, self.mp_group) |
| 673 | else: |
| 674 | continue |
| 675 | if len(child._buffers) != 0 and self.state_dict is not None: |
| 676 | Loading.load_buffer(child, self.state_dict, checking_key) |
| 677 | |
| 678 | # When using partition_config (custom patterns/presets), use pattern-based routing |
| 679 | # instead of linear_policies. This keeps all pattern logic centralized here. |
| 680 | if self.partition_config is not None: |
| 681 | full_name = class_name + '.' + name if class_name else name |
| 682 | if full_name in self._tied_gathered_column_module_names: |
| 683 | continue |
| 684 | elif isinstance(child, nn.Embedding): |
| 685 | # Check if embedding matches any pattern |
| 686 | param_name = full_name + ".weight" |
| 687 | model_type = self._get_model_type() |
| 688 | spec = self.partition_config.find_matching_spec(param_name, model_type) |
| 689 | if spec is not None and spec.partition_type != PartitionType.SKIP: |
| 690 | new_child = self._slice_embedding(child, full_name, False) |
| 691 | if new_child is not None: |
| 692 | setattr(r_module, name, new_child) |
| 693 | # If no pattern matched or skip, leave embedding unchanged |
| 694 | elif hasattr(child, "weight") and getattr(child.weight, "dim", lambda: 0)() == 2: |
| 695 | new_child = self._replace_with_config(child, full_name) |
| 696 | if new_child is not None: |
| 697 | setattr(r_module, name, new_child) |
| 698 | else: |
| 699 | self.update_mp_params(child) |
| 700 | self._replace_module(child, name, class_name) |
| 701 | # Traditional path: use linear_policies for type-based routing |
| 702 | elif child.__class__ in self.linear_policies: |
| 703 | setattr(r_module, name, self.linear_policies[child.__class__](child, prev_name + '.' + name, |
| 704 | self.conv_linear_layer)) |
| 705 | elif any(isinstance(child, lp) for lp in self.linear_policies): |
| 706 | # Added for falcon model support |
| 707 | # Note: isinstance will account for class inheritance, child.__class__ does not |
| 708 | key = None |
| 709 | for lp in self.linear_policies: |
| 710 | if isinstance(child, lp): |