@@ -276,11 +276,11 @@ Below is an example of a YAML template for replacing all original Linear modules
...
@@ -276,11 +276,11 @@ Below is an example of a YAML template for replacing all original Linear modules
name:"^model\\.layers\\..*$"# regular expression
name:"^model\\.layers\\..*$"# regular expression
class:torch.nn.Linear# only match modules matching name and class simultaneously
class:torch.nn.Linear# only match modules matching name and class simultaneously
replace:
replace:
class:ktransformers.operators.linear.KTransformerLinear# optimized Kernel on quantized data types
class:ktransformers.operators.linear.KTransformersLinear# optimized Kernel on quantized data types
device:"cpu"# which devices to load this module when initializing
device:"cpu"# which devices to load this module when initializing
kwargs:
kwargs:
generate_device:"cuda"
generate_device:"cuda"
generate_linear_type:"QuantizedLinearMarlin"
generate_linear_type:"KLinearMarlin"
```
```
Each rule in the YAML file has two parts: `match` and `replace`. The `match` part specifies which module should be replaced, and the `replace` part specifies the module to be injected into the model along with the initialization keywords.
Each rule in the YAML file has two parts: `match` and `replace`. The `match` part specifies which module should be replaced, and the `replace` part specifies the module to be injected into the model along with the initialization keywords.
As we can see, each rule in the YAML file has two parts: `match` and `replace`.
As we can see, each rule in the YAML file has two parts: `match` and `replace`.
...
@@ -98,9 +98,9 @@ The match part specifies which module should be replaced, and the replace part s
...
@@ -98,9 +98,9 @@ The match part specifies which module should be replaced, and the replace part s
<h3id="experts">Routed Experts </h3>
<h3id="experts">Routed Experts </h3>
For routed experts, the module we inject is a wrapper of CPUInfer, KTransformersMLPExpert. There are several implementations within a wrapper, and we need to specify keywords to tell the wrapper which implementation we want to use and how we intend to use it.
For routed experts, the module we inject is a wrapper of CPUInfer, KTransformersExperts. There are several implementations within a wrapper, and we need to specify keywords to tell the wrapper which implementation we want to use and how we intend to use it.
In KTransformers, some models exhibit different behaviors during prefilling and generation for better performance. KTransformersMLPExpert is one of them. All these special modules have a `device` keyword describing which device the module should be initialized on. Other keywords specify the behaviors during prefilling and generation and may be differ when using different injection modules. Here, we specify which implementation on which device we want to use during prefilling and generation, and which device the output should be on.
In KTransformers, some models exhibit different behaviors during prefilling and generation for better performance. KTransformersExperts is one of them. All these special modules have a `device` keyword describing which device the module should be initialized on. Other keywords specify the behaviors during prefilling and generation and may be differ when using different injection modules. Here, we specify which implementation on which device we want to use during prefilling and generation, and which device the output should be on.
Note that we only use these parameters when layer-wise prefilling is enabled; otherwise, prefilling is conducted with the same configuration as generation.
Note that we only use these parameters when layer-wise prefilling is enabled; otherwise, prefilling is conducted with the same configuration as generation.
In the original implementation of Transformers, MoE is implemented using `nn.ModuleList`. We don't want KTransformers to iterate through all the sub-modules in the list, so we set `recursive: False` in this rule to prevent recursive injection into submodules of the current module. Here is the YAML rule:
In the original implementation of Transformers, MoE is implemented using `nn.ModuleList`. We don't want KTransformers to iterate through all the sub-modules in the list, so we set `recursive: False` in this rule to prevent recursive injection into submodules of the current module. Here is the YAML rule:
...
@@ -109,13 +109,13 @@ In the original implementation of Transformers, MoE is implemented using `nn.Mod
...
@@ -109,13 +109,13 @@ In the original implementation of Transformers, MoE is implemented using `nn.Mod
-match:
-match:
name:"^model\\.layers\\..*\\.mlp\\.experts$"
name:"^model\\.layers\\..*\\.mlp\\.experts$"
replace:
replace:
class:ktransformers.operators.experts.KTransformersMLPExpert# custom MoE Kernel with expert parallelism
class:ktransformers.operators.experts.KTransformersExperts# custom MoE Kernel with expert parallelism
device:"cpu"# device to load this module on initialization
device:"cpu"# device to load this module on initialization
kwargs:
kwargs:
prefill_device:"cuda"
prefill_device:"cuda"
prefill_mlp_type:"MLPExpertsTorch"
prefill_op:"KExpertsTorch"
generate_device:"cpu"
generate_device:"cpu"
generate_mlp_type:"MLPCPUExperts"
generate_op:"KExpertsCPU"
out_device:"cuda"
out_device:"cuda"
recursive:False# don't recursively inject submodules of this module
recursive:False# don't recursively inject submodules of this module
```
```
...
@@ -126,7 +126,7 @@ If we inject the expert list as a custom module, we can't use the interface in `
...
@@ -126,7 +126,7 @@ If we inject the expert list as a custom module, we can't use the interface in `
print(f"This linear module's in_features or out_features is not divisible by GPTQ_MARLIN_MIN_THREAD_N({GPTQ_MARLIN_MIN_THREAD_N}), using QuantizedLinearTorch instead.")
print(f"This linear module's in_features or out_features is not divisible by GPTQ_MARLIN_MIN_THREAD_N({GPTQ_MARLIN_MIN_THREAD_N}), using KLinearTorch instead.")
print(f"This linear module's in_features or out_features is not divisible by GPTQ_MARLIN_MIN_THREAD_N({GPTQ_MARLIN_MIN_THREAD_N}), using QuantizedLinearTorch instead.")
print(f"This linear module's in_features or out_features is not divisible by GPTQ_MARLIN_MIN_THREAD_N({GPTQ_MARLIN_MIN_THREAD_N}), using KLinearTorch instead.")
