feat(CRITICAL): 从 GitHub 扫描搬运 ixformer SDK + xllm 完整 GDN/MoE 代码

来源:
  1. Chranos/ixformer (GitHub) → ixformer_sdk/ (230 files, 70K lines)
     - inference/functions/vllm.py: vllm_moe_topk_softmax 完整实现 (2033 lines)
     - inference/functions/moe.py: MoE ops 完整实现 (1380 lines)
     - contrib/vllm_flash_attn/: FA2 Python 接口 (1018 lines)
     - contrib/tgi/fused_moe.py: TGI fused MoE (429 lines)
     - csrc/include/ixformer/: C++ kernel headers + cmake

  2. Deep-Spark/xllm (GitHub) → upstream_ref/xllm_latest/ (+15 files)
     - npu_torch/qwen3_5_decoder_layer_impl.cpp/.h
     - npu_torch/qwen3_5_gated_delta_net.cpp/.h
     - npu_torch/qwen3_next_*.cpp/.h (6 files)
     - npu_torch/attention.cpp/.h + fused_moe.cpp/.h + CMakeLists.txt
     - models/llm/qwen3_5.h + qwen3_5_mtp.h + qwen3_next.h
     - models/vlm/qwen3_5.h

调用链完整性:
  ixformer_sdk/inference/functions/vllm.py
    → ops.infer.moe_topk_softmax() (C++ 层)
    → 这就是 base 镜像 libixformer.so 里的实现

  upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
    → ixformer::infer::topk_softmax() (直接 C++ 调用)
    → ixformer::infer::group_gemm() → 完整 7-step MoE pipeline

ixformer_sdk/train/speedformer/policy/baichuan.py

@@ -0,0 +1,59 @@
+import warnings
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Callable, Dict, List, Union
+
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import Module
+
+from ixformer.train.speedformer.policy.utils import SubModuleReplacementDescription
+from ixformer.train.speedformer.policy.replacer import Replacer
+
+from ixformer.train.speedformer.models.baichuan.modeling_baichuan import BaichuanModel, DecoderLayer
+from ixformer.train.speedformer.layers.normalization import APEXFusedRMSNorm, IXFFusedRMSNorm
+from ixformer.train.speedformer.layers.baichuan.attention import BaichuanAttention
+from ixformer.train.speedformer.layers.baichuan.mlp import IXFBaichuanMLP
+
+
+class BaichuanReplacer(Replacer):
+    def __init__(self):
+        self.policy = {}
+
+    def module_policy(self) -> Dict[Union[str, nn.Module], List[SubModuleReplacementDescription]]:
+        self.append_or_create_submodule_replacement(
+            description=[
+                SubModuleReplacementDescription(
+                    suffix="input_layernorm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={}
+                ),
+                SubModuleReplacementDescription(
+                    suffix="post_attention_layernorm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={},
+                ),
+                SubModuleReplacementDescription(
+                    suffix="self_attn",
+                    target_module=BaichuanAttention,
+                    kwargs={}
+                ),
+                SubModuleReplacementDescription(
+                    suffix="mlp",
+                    target_module=IXFBaichuanMLP,
+                    kwargs={}
+                ),
+            ],
+            target_key="DecoderLayer"
+        )
+
+        self.append_or_create_submodule_replacement(
+            description=[
+                SubModuleReplacementDescription(
+                    suffix="norm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={}
+                ),
+            ],
+            target_key=BaichuanModel
+        )

ixformer_sdk/train/speedformer/policy/bloom.py

@@ -0,0 +1,53 @@
+import warnings
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Callable, Dict, List, Union
+
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import Module
+
+from ixformer.train.speedformer.policy.utils import SubModuleReplacementDescription
+from ixformer.train.speedformer.policy.replacer import Replacer
+
+from ixformer.train.speedformer.models.bloom.modeling_bloom import BloomModel, BloomBlock
+from ixformer.train.speedformer.layers.normalization import APEXFusedRMSNorm, IXFFusedRMSNorm
+from ixformer.train.speedformer.layers.bloom.attention import BloomFlashAttention
+
+
+class BloomReplacer(Replacer):
+    def __init__(self):
+        self.policy = {}
+
+    def module_policy(self) -> Dict[Union[str, nn.Module], List[SubModuleReplacementDescription]]:
+        self.append_or_create_submodule_replacement(
+            description=[
+                SubModuleReplacementDescription(
+                    suffix="input_layernorm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={}
+                ),
+                SubModuleReplacementDescription(
+                    suffix="post_attention_layernorm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={},
+                ),
+                SubModuleReplacementDescription(
+                    suffix="self_attention",
+                    target_module=BloomFlashAttention,
+                    kwargs={}
+                ),
+            ],
+            target_key="BloomBlock"
+        )
+
+        self.append_or_create_submodule_replacement(
+            description=[
+                SubModuleReplacementDescription(
+                    suffix="ln_f",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={}
+                ),
+            ],
+            target_key=BloomModel
+        )

ixformer_sdk/train/speedformer/policy/chatglm.py

@@ -0,0 +1,57 @@
+from typing import Callable, Dict, List, Union
+from torch.nn import Module
+
+from ixformer.train.speedformer.policy.utils import SubModuleReplacementDescription
+from ixformer.train.speedformer.policy.replacer import Replacer
+
+from ixformer.train.speedformer.layers.normalization import APEXFusedRMSNorm, IXFFusedRMSNorm
+from ixformer.train.speedformer.layers.chatglm.attention import ChatglmFlashAttention
+from ixformer.train.speedformer.layers.chatglm.methods import ChatGLMModel_forward
+
+
+class ChatglmReplacer(Replacer):
+    def __init__(self):
+        self.policy = {}
+
+    def module_policy(self) -> Dict[str | Module, List[SubModuleReplacementDescription]]:
+        self.append_or_create_submodule_replacement(
+            description=[
+                SubModuleReplacementDescription(
+                    suffix="final_layernorm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={}
+                ),
+            ],
+            target_key="GLMTransformer"
+        )
+        self.append_or_create_submodule_replacement(
+            description=[
+                SubModuleReplacementDescription(
+                    suffix="input_layernorm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={}
+                ),
+                SubModuleReplacementDescription(
+                    suffix="post_attention_layernorm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={}
+                ),
+            ],
+            target_key="GLMBlock"
+        )
+        self.append_or_create_submodule_replacement(
+            description=[
+                SubModuleReplacementDescription(
+                    suffix="self_attention",
+                    target_module=ChatglmFlashAttention,
+                    kwargs={}
+                ),
+            ],
+            target_key="GLMBlock"
+        )
+        self.append_or_create_method_replacement(
+            description=[
+                {"forward": ChatGLMModel_forward()}
+            ],
+            target_key="ChatGLMModel"
+        )

ixformer_sdk/train/speedformer/policy/gpt2.py

@@ -0,0 +1,27 @@
+import torch
+import torch.nn as nn
+from torch.nn import LayerNorm
+from types import ModuleType, MethodType
+from abc import ABC
+
+from ixformer.train.speedformer.models.gpt2.modeling_gpt2 import GPT2FlashAttention2
+
+from ixformer.train.speedformer.layers.normalization import replace_layernorm_forward
+from ixformer.train.speedformer.layers.gpt2.attention import replace_flash_attn_forward
+
+
+class GPT2Replacer(ABC):
+    def __init__(self) -> None:
+        super().__init__()
+
+    @staticmethod
+    def accelerate(model):
+        # layer/kernel replace
+        for name, module in model.named_modules():
+            if isinstance(module, LayerNorm):
+                module.forward = MethodType(replace_layernorm_forward, module)
+            if isinstance(module, GPT2FlashAttention2):
+                module._flash_attention_forward = MethodType(
+                    replace_flash_attn_forward, module)
+
+        return model

ixformer_sdk/train/speedformer/policy/llama.py

@@ -0,0 +1,104 @@
+import warnings
+import types
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Callable, Dict, List, Union
+
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import Module
+
+from ixformer.train.speedformer.policy.utils import SubModuleReplacementDescription
+from ixformer.train.speedformer.policy.replacer import Replacer
+
+from ixformer.train.speedformer.layers.normalization import APEXFusedRMSNorm, IXFFusedRMSNorm
+from ixformer.train.speedformer.layers.llama.attention import LlamaAttention as IXF_LlamaAttention
+from ixformer.train.speedformer.layers.llama.mlp import IXFLlamaMLP
+from ixformer.train.speedformer.layers.llama.llama_method import LlamaModel_forward, LlamaForCausalLM_forward
+from ixformer.train.speedformer.layers.fast_lora.fast_lora import apply_lora_mlp_swiglu
+
+from peft import PeftType
+
+
+class LlamaReplacer(Replacer):
+    def __init__(self):
+        self.policy = {}
+
+    def module_policy(self) -> Dict[Union[str, nn.Module], List[SubModuleReplacementDescription]]:
+        self.append_or_create_submodule_replacement(
+            description=[
+                SubModuleReplacementDescription(
+                    suffix="input_layernorm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={}
+                ),
+                SubModuleReplacementDescription(
+                    suffix="post_attention_layernorm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={},
+                ),
+                SubModuleReplacementDescription(
+                    suffix="self_attn",
+                    target_module=IXF_LlamaAttention,
+                    kwargs={}
+                ),
+                # SubModuleReplacementDescription(
+                #     suffix="mlp",
+                #     target_module=IXFLlamaMLP,
+                #     kwargs={}
+                # ),
+            ],
+            target_key="LlamaDecoderLayer"
+        )
+
+        self.append_or_create_submodule_replacement(
+            description=[
+                SubModuleReplacementDescription(
+                    suffix="norm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={}
+                ),
+            ],
+            target_key="LlamaModel"
+        )
+
+        self.append_or_create_method_replacement(
+            description=[
+                {"forward": LlamaModel_forward()}
+            ],
+            target_key="LlamaModel"
+        )
+        self.append_or_create_method_replacement(
+            description=[
+                {"forward": LlamaForCausalLM_forward()}
+            ],
+            target_key="LlamaForCausalLM"
+        )
+
+    def post_process(self, model: nn.Module):
+        if model.peft_type != PeftType.LORA:
+            return
+        peft_config = model.peft_config
+        active_adapter = model.active_adapters[0] if \
+            hasattr(model, "active_adapters") else model.active_adapter
+        target_modules = peft_config[active_adapter].target_modules
+
+        # for now, fast_lora only support lora_dropout=0 and bias=None
+        lora_dropout = model.peft_config[active_adapter].lora_dropout
+        bias = model.peft_config[active_adapter].bias
+
+        # 首先判断是否可以使用fast_lora
+        check = lora_dropout == 0 and bias == "none"
+
+        # 其次确定mlp的3个线性层是否在target_modules
+        mlp_use_fastlora = "gate_proj" in target_modules and "up_proj" in target_modules and "up_proj" in target_modules
+
+        n_mlp = 0
+        if check:
+            if mlp_use_fastlora:
+                for layer in model.model.model.layers:
+                    layer.mlp.forward = types.MethodType(
+                        apply_lora_mlp_swiglu, layer.mlp)
+                    n_mlp += 1
+
+        print(f"{len(model.model.model.layers)} layers replace mlp with fast_lora mlp")

ixformer_sdk/train/speedformer/policy/qwen2.py

@@ -0,0 +1,57 @@
+import warnings
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Callable, Dict, List, Union
+
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import Module
+
+from ixformer.train.speedformer.policy.utils import SubModuleReplacementDescription
+from ixformer.train.speedformer.policy.replacer import Replacer
+import os
+import sys
+from ixformer.train.speedformer.layers.normalization import APEXFusedRMSNorm, IXFFusedRMSNorm
+from ixformer.train.speedformer.layers.qwen2.attention import QwenAttention as IXF_QwenAttention
+
+
+class Qwen2Replacer(Replacer):
+    def __init__(self):
+        self.policy = {}
+        
+    def module_policy(self) -> Dict[Union[str, nn.Module], List[SubModuleReplacementDescription]]:
+        self.append_or_create_submodule_replacement(
+            description=[
+                SubModuleReplacementDescription(
+                    suffix="input_layernorm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={}
+                ),
+                SubModuleReplacementDescription(
+                    suffix="post_attention_layernorm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={},
+                ),
+                SubModuleReplacementDescription(
+                    suffix="self_attn",
+                    target_module=IXF_QwenAttention,
+                    kwargs={}
+                ),
+            ],
+            target_key="Qwen2DecoderLayer"
+        )
+        
+        self.append_or_create_submodule_replacement(
+            description=[
+                SubModuleReplacementDescription(
+                    suffix="norm",
+                    target_module=APEXFusedRMSNorm,
+                    kwargs={}
+                ),
+            ],
+            target_key="Qwen2Model"
+        )
+        
+
+        
+    

ixformer_sdk/train/speedformer/policy/replacer.py

@@ -0,0 +1,224 @@
+import warnings
+from types import MethodType
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Any, Callable, Dict, List, Optional, Set, Union
+import tabulate
+
+import torch.nn as nn
+
+from ixformer.train.speedformer.policy.utils import SubModuleReplacementDescription, ModulePolicyDescription, getattr_, setattr_, print_rank_0
+
+
+class Replacer(ABC):
+    def __init__(self):
+        self.policy = {}
+
+    def module_policy(self) -> Dict[Union[str, nn.Module], List[SubModuleReplacementDescription]]:
+        r"""
+        This method returns the module policy, which is a dictionary. The key is the module name or the module object,
+        and the value is the ModulePolicyDescription object. The ModulePolicyDescription object describes how the module
+        will be transformed.
+        """
+
+    def append_or_create_submodule_replacement(
+        self,
+        description: Union[SubModuleReplacementDescription, List[SubModuleReplacementDescription]],
+        target_key: Union[str, nn.Module],
+    ) -> Dict[Union[str, nn.Module], List]:
+        r"""
+        Append or create a new submodule replacement description to the policy for the given key.
+
+        Args:
+            submodule_replace_desc (Union[SubModuleReplacementDescription, List[SubModuleReplacementDescription]]): the submodule replacement description to be appended
+            policy (Dict[Union[str, nn.Module], ModulePolicyDescription]): the policy to be updated
+            target_key (Union[str, nn.Module]): the key of the policy to be updated
+        """
+        # convert to list
+        if isinstance(description, SubModuleReplacementDescription):
+            description = [description]
+
+        # append or create a new description
+        if target_key in self.policy:
+            if self.policy[target_key].sub_module_replacement is None:
+                self.policy[target_key].sub_module_replacement = description
+            else:
+                self.policy[target_key].sub_module_replacement.extend(
+                    description)
+        else:
+            self.policy[target_key] = ModulePolicyDescription(
+                sub_module_replacement=description)
+
+    def append_or_create_method_replacement(
+        self,
+        description: Dict[str, Callable],
+        target_key: Union[str, nn.Module],
+    ) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
+        r"""
+        Append or create a new method replacement description to the policy for the given key.
+
+        Args:
+            description (Union[SubModuleReplacementDescription, List[SubModuleReplacementDescription]]): the submodule replacement description to be appended
+            policy (Dict[Union[str, nn.Module], ModulePolicyDescription]): the policy to be updated
+            target_key (Union[str, nn.Module]): the key of the policy to be updated
+        """
+        if target_key in self.policy:
+            if self.policy[target_key].method_replacement is None:
+                self.policy[target_key].method_replacement = description
+            else:
+                self.policy[target_key].method_replacement.extend(description)
+        else:
+            self.policy[target_key] = ModulePolicyDescription(
+                method_replacement=description)
+
+    def append_or_create_attribute_replacement(
+        self,
+        description: Dict[str, Callable],
+        target_key: Union[str, nn.Module],
+    ) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
+        r"""
+        Append or create a new method replacement description to the policy for the given key.
+
+        Args:
+            description (Union[SubModuleReplacementDescription, List[SubModuleReplacementDescription]]): the submodule replacement description to be appended
+            policy (Dict[Union[str, nn.Module], ModulePolicyDescription]): the policy to be updated
+            target_key (Union[str, nn.Module]): the key of the policy to be updated
+        """
+        if target_key in self.policy:
+            if self.policy[target_key].attribute_replacement is None:
+                self.policy[target_key].attribute_replacement = description
+            else:
+                self.policy[target_key].attribute_replacement.extend(
+                    description)
+        else:
+            self.policy[target_key] = ModulePolicyDescription(
+                attribute_replacement=description)
+
+    def accelerate(self, model) -> None:
+        r"""
+        Replace the module according to the policy, and replace the module one by one
+
+        Args:
+            model (:class:`torch.nn.Module`): The model to shard
+        """
+        self.module_policy()
+        self.module_replace = []
+        for layer_cls, module_description in self.policy.items():
+            self.replace_sub_module(
+                model, layer_cls, module_description.sub_module_replacement)
+            self._replace_method(
+                model, layer_cls, module_description.method_replacement)
+        print_rank_0(tabulate.tabulate(self.module_replace, headers=[
+                     "old_layer", "new_layer"], tablefmt="psql"))
+        return model
+
+    def replace_sub_module(
+        self,
+        module: nn.Module,
+        origin_cls: Union[str, nn.Module],
+        sub_module_replacement: List[SubModuleReplacementDescription],
+    ) -> None:
+        r"""
+        Reverse the replace layer operation
+        """
+        if not sub_module_replacement:
+            return
+
+        if (isinstance(origin_cls, str) and origin_cls == module.__class__.__name__) or (
+            module.__class__ == origin_cls
+        ):
+            for description in sub_module_replacement:
+                suffix = description.suffix
+                target_module = description.target_module
+                kwargs = {} if description.kwargs is None else description.kwargs
+
+                assert target_module is not None, "target_module should not be None"
+
+                native_sub_module = getattr_(module, suffix, ignore=True)
+
+                assert not isinstance(
+                    native_sub_module, target_module
+                ), f"The module with suffix {suffix} has been replaced, please check the policy"
+
+                # if it is None and we are allowed to ignore this module
+                # just skip
+                if description.ignore_if_not_exist and native_sub_module is None:
+                    continue
+                try:
+                    replace_layer = target_module.from_native_module(
+                        native_sub_module, **kwargs)
+                except Exception as e:
+                    raise RuntimeError(
+                        f"Failed to replace {suffix} of type {native_sub_module.__class__.__qualname__}"
+                        f" with {target_module.__qualname__} with the exception: {e}. "
+                        "Please check your model configuration or sharding policy, you can set up an issue for us to help you as well."
+                    )
+
+                setattr_(module, suffix, replace_layer)
+                self.module_replace.append(
+                    [native_sub_module.__class__.__qualname__, target_module.__qualname__])
+
+        for name, child in module.named_children():
+            self.replace_sub_module(
+                child,
+                origin_cls,
+                sub_module_replacement,
+            )
+
+    def _replace_method(self, module: nn.Module, origin_cls: Union[str, nn.Module], method_replacement: List[Dict[str, Callable]]):
+        if not method_replacement:
+            return
+
+        if (isinstance(origin_cls, str) and origin_cls == module.__class__.__name__) or (
+            module.__class__ == origin_cls
+        ):
+            for method in method_replacement:
+                for method_name, new_method in method.items():
+                    # bind the new method to the module
+                    bound_method = MethodType(new_method, module)
+                    setattr(module, method_name, bound_method)
+
+        for name, child in module.named_children():
+            self._replace_method(
+                child,
+                origin_cls,
+                method_replacement,
+            )
+
+    def _replace_attr(
+        self,
+        module: nn.Module,
+        origin_cls: Union[str, nn.Module],
+        attr_replacement: List[Dict[str, Any]],
+    ) -> None:
+        r"""
+        Replace the attribute of the layer
+
+        Args:
+            module (:class:`torch.nn.Module`): The object of layer to shard
+            attr_replacement (Dict): The attribute dict to modify
+        """
+        if not attr_replacement:
+            return
+
+        if (isinstance(origin_cls, str) and origin_cls == module.__class__.__name__) or (
+            module.__class__ == origin_cls
+        ):
+            for attr in attr_replacement:
+                for module_attr, target_attr in attr.items():
+                    native_attr = getattr_(module, module_attr, ignore=False)
+                    if isinstance(native_attr, type):
+                        replace_attr = target_attr.from_native_attr(
+                            native_attr)
+                        setattr_(module, module_attr,
+                                 replace_attr, ignore=False)
+                    else:
+                        setattr_(module, module_attr,
+                                 target_attr, ignore=False)
+
+        for name, child in module.named_children():
+            self._replace_attr(
+                child,
+                origin_cls,
+                attr_replacement,
+            )

ixformer_sdk/train/speedformer/policy/utils.py

@@ -0,0 +1,156 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Union
+import re
+import torch
+import torch.nn as nn
+
+
+@dataclass
+class SubModuleReplacementDescription:
+    r"""
+    Describe how a submodule will be replaced
+
+    Args:
+        suffix (str): used to get the submodule object
+        target_module (ParallelModule): specifies the module class used to replace to submodule
+        kwargs (Dict[str, Any]): the dictionary used to pass extra arguments to the `ParallelModule.from_native_module` method.
+        ignore_if_not_exist (bool): if the submodule does not exist, ignore it or raise an exception
+    """
+
+    suffix: str
+    target_module: nn.Module
+    kwargs: Dict[str, Any] = None
+    ignore_if_not_exist: bool = False
+
+
+@dataclass
+class ModulePolicyDescription:
+    "copy from colossalai, for now sub_module_replacement and method_replacement is used"
+    r"""
+    Describe how the attributes and parameters will be transformed in a policy.
+
+    Args:
+        attribute_replacement (Dict[str, Any]): key is the attribute name, value is the attribute value after sharding
+        param_replacement (List[Callable]): a list of functions to perform in-place param replacement. The function
+                    must receive only one arguments: module. One example is
+
+                    ```python
+                    def example_replace_weight(module: torch.nn.Module):
+                        weight = module.weight
+                        new_weight = shard_rowwise(weight, process_group)
+                        module.weight = torch.nn.Parameter(new_weight)
+                    ```
+        sub_module_replacement (List[SubModuleReplacementDescription]): each element in the list is a SubModuleReplacementDescription
+                    object which specifies the module to be replaced and the target module used to replacement.
+        method_replace (Dict[str, Callable]): key is the method name, value is the method for replacement
+    """
+
+    attribute_replacement: List[Dict[str, Any]] = None
+    param_replacement: List[Callable] = None
+    sub_module_replacement: List[SubModuleReplacementDescription] = None
+    method_replacement: List[Dict[str, Callable]] = None
+
+
+def getattr_(obj, attr: str, ignore: bool = False):
+    r"""
+    Get the object's multi sublevel attr
+
+    Args:
+        obj (object): The object to set
+        attr (str): The multi level attr to set
+        ignore (bool): Whether to ignore when the attr doesn't exist
+    """
+
+    attrs = attr.split(".")
+    for a in attrs:
+        try:
+            obj = get_obj_list_element(obj, a)
+        except AttributeError:
+            if ignore:
+                return None
+            raise AttributeError(
+                f"Object {obj.__class__.__name__} has no attribute {attr}")
+    return obj
+
+
+def get_obj_list_element(obj, attr: str):
+    r"""
+    Get the element of the list in the object
+
+    If the attr is a normal attribute, return the attribute of the object.
+    If the attr is a index type, return the element of the index in the list, like `layers[0]`.
+
+    Args:
+        obj (Object): The object to get
+        attr (str): The suffix of the attribute to get
+
+    """
+    re_pattern = r"\[\d+\]"
+    prog = re.compile(re_pattern)
+    result = prog.search(attr)
+    if result:
+        matched_brackets = result.group()
+        matched_index = matched_brackets.replace("[", "")
+        matched_index = matched_index.replace("]", "")
+        attr_ = attr.replace(matched_brackets, "")
+        container_obj = getattr(obj, attr_)
+        obj = container_obj[int(matched_index)]
+    else:
+        obj = getattr(obj, attr)
+    return obj
+
+
+def setattr_(obj, attr: str, value, ignore: bool = False):
+    r"""
+    Set the object's multi sublevel attr to value, if ignore, ignore when it doesn't exist
+
+    Args:
+        obj (object): The object to set
+        attr (str): The multi level attr to set
+        value (Any): The value to set
+        ignore (bool): Whether to ignore when the attr doesn't exist
+    """
+
+    attrs = attr.split(".")
+    for a in attrs[:-1]:
+        try:
+            obj = get_obj_list_element(obj, a)
+        except AttributeError:
+            if ignore:
+                return
+            raise AttributeError(
+                f"Object {obj.__class__.__name__} has no attribute {attr}")
+    set_obj_list_element(obj, attrs[-1], value)
+
+
+def set_obj_list_element(obj, attr: str, value):
+    r"""
+    Set the element to value of a list object
+
+    It used like set_obj_list_element(obj, 'layers[0]', new_layer), it will set obj.layers[0] to value
+
+    Args:
+        obj (object): The object to set
+        attr (str): the string including a list index like `layers[0]`
+    """
+    re_pattern = r"\[\d+\]"
+    prog = re.compile(re_pattern)
+    result = prog.search(attr)
+    if result:
+        matched_brackets = result.group()
+        matched_index = matched_brackets.replace("[", "")
+        matched_index = matched_index.replace("]", "")
+        attr_ = attr.replace(matched_brackets, "")
+        container_obj = getattr(obj, attr_)
+        container_obj[int(matched_index)] = value
+    else:
+        setattr(obj, attr, value)
+
+
+def print_rank_0(message):
+    if torch.distributed.is_initialized():
+        if torch.distributed.get_rank() == 0:
+            print(message, flush=True)
+    else:
+        print(message, flush=True)