Repository: OpenBMB/BMCook (language: python)

Issue:
Sprune dev



Discussion / hints:

Fix patch:
diff --git a/README-ZH.md b/README-ZH.md
index d792417..a70fbbd 100644
--- a/README-ZH.md
+++ b/README-ZH.md
@@ -129,15 +129,15 @@ $ git clone git@github.com:OpenBMB/BMCook.git
 
 ## 压缩效果
 
-基于GPT-J (6B)，我们评估了不同的压缩组合，语料库使用了OpenWebText。同时，我们基于该语料从头训练了一个0.7B参数的GPT-J作为对比。
+基于T5-3B，我们评估了不同的压缩组合，压缩语料库使用了Pile。选择SST-2、MNLI和SQuAD作为下游任务进行评测。适配下游任务时，我们固定了预训练模型参数，采用adapter-tuning进行训练。
 
-|                        |     LM Loss    |     Relative Performance    |     Speedup    |
+|                        |     Average Performance    |     Relative Performance    |     Speedup    |
 |------------------------|----------------|-----------------------------|----------------|
-|     GPT-J              |           3.37 |                        -    |          1x    |
-|     GPT-J (0.7B)       |           4.06 |                       83.0% |         ~10x   |
-|     GPT-J (P+D)        |           3.57 |                       94.4% |          2x    |
-|     GPT-J (P+D+Q)      |           3.58 |                       94.1% |          8x    |
-|     GPT-J (P+D+Q+M)    |           3.69 |                       91.3% |          10x   |
+|     T5-3B              |           0.9258 |                        -    |          1x    |
+|     T5-Base       |           0.8796 |                       95.0% |         7x   |
+|     T5-3B (P+D)        |           0.9150 |                       98.8% |          2x    |
+|     T5-3B (P+D+Q)      |           0.9126 |                       98.6% |          8x    |
+|     T5-3B (P+D+Q+M)    |           0.9017 |                       97.4% |          12x   |
 
 D 代表知识蒸馏；P 代表模型剪枝；Q 代表模型量化；M 代表模型专家化。
 
diff --git a/README.md b/README.md
index e5f9393..f357b12 100644
--- a/README.md
+++ b/README.md
@@ -9,7 +9,7 @@
 
 
 <p align="center">
-  <a href="#overview">Overview</a> • <a href="#documentation">Documentation</a> • <a href="#install">Installation</a> • <a href="#quick-start">Quick Start</a> • <a href="./README-ZH.md" target="_blank">简体中文</a>
+  <a href="#overview">Overview</a> • <a href="#documentation">Documentation</a> • <a href="#install">Installation</a> • <a href="#usage">Usage</a> • <a href="#quick-start">Quick Start</a> • <a href="./README-ZH.md" target="_blank">简体中文</a>
 <br>
 </p>
 
@@ -67,20 +67,128 @@ $ python3 setup.py install
 
 Please refer to [the installation guide](https://bmtrain.readthedocs.io/en/latest/) of BMTrain for more details.
 
-Then, clone the repository.
+Then, install BMCook. 
 
+**From PyPI (Recommend)**
+
+```shell
+$ pip install bmcook
+```
+
+**From source**
 
 ```shell
 $ git clone git@github.com:OpenBMB/BMCook.git
+cd BMCook
+python3 setup.py install
+```
+
+<div id="usage"></div>
+
+## Usage
+
+### 1. Design your BMCook config.
+You should give a json file to state your compress strategy.
+
+```json
+{ "distillation": {
+    "ce_scale": 0,
+    "ce_temp": 1,
+      
+    "mse_hidn_scale": 0,
+    "mse_hidn_module": ['[placehold]'],
+    "mse_hidn_proj": false,
+      
+    "mse_att_scale": 0,
+    "mse_att_module": ['[placehold]'],
+  },
+
+  "pruning": {
+    "is_pruning": false,
+    "pruning_mask_path": None,
+    "pruned_module": ['[placehold]'],
+    "mask_method": "m4n2_1d/m4n2_2d/sprune",
+    "sprune": {
+        "criterion": "l0",
+        "training_mask": ['[placehold]'],
+        "fixed_mask_path": "",
+        "mask_mode": "train_mask",
+        "target_sparsity": 0.5
+    }
+  },
+
+  "quantization": {
+    "is_quant": false,
+    "quantized_module": [],
+  },
+
+  "MoEfication": {
+    "is_moefy": false,
+    "first_FFN_module": ['[placehold]'],
+  }
+}
 ```
+To notice:
+
+- `is_moefy`, `is_quant`, `ispruning` are switch parameters. If false, other parameters will be blocked. `mask_method` takes similar works. When `mask_method` is "m4n2_1d" or "m4n2_2d", it will execute unstructure pruning, but when is "sprune", the `sprune` field will be activated. For distillation, when the `ce_scale` or `mse_hidn_scale` is greater than 0, the corresponding distilling mode will be switched on.
+
+- It's not recommended to use MoE and Distilling simultaneously.
+
+### 2. Basic usage in your code.
+BMCook provides unified interface `CookTrainer`. BMCook will introduce distillation pruning and MoEfication, which may add some terms to model outputs. You can use it to manage your model, and these modifications.
+```python
+from bmcook import CookTrainer
+from bmcook.utils.config import ConfigParser
+
+#prepare your model, dataloader and optimizer...
+...
+
+# setting up your BMCook strategy
+CookTrainer.set_compression(cookconfig, model, optimizer, model_distill)
+
+# train
+for data in dataloader:
+    targets = ...
+    ...
+    outputs = CookTrainer.forward(model, loss_func, targets, *your_model_inputs, **your_model_kwinputs)
+
+    [loss, model_outputs, lag_loss, sparsity, distill_loss] = outputs
+```
+the loss equals to the sum of model_loss, lag_loss and distill_loss. So if you wanna know the model performance, please minus them. Noticed that if sprune is not setted, the lag_loss and loss_func will be `None`, so do distilling.
+```python
+model_loss = loss - lag_loss - distill_loss # sprune and distilling both setted.
+model_loss = loss - distill_loss # only distilling used. 
+```
+
+BMCook also provides discrete interfaces to initialize compression settings. If you want to design your Trainer for your own needs, you can use these discrete interfaces. Noticed that the output format should keep the same with `CookTrainer` when you define your own Trainer. For details about extension on `CookTrainer`, you can refer to `CPMAntTrainer`.
+```python
+from bmcook import BMDistill
+
+# Define your own Trainer. 
+Trainer = ...
+
+# Set up the distillation
+Trainer.forward = BMDistill.set_forward(model, teacher, Trainer.forward, cook_config)
+```
+
+### 3. How to run your code
+You can run your code as normal, but should state where your cookconfig is:
+```shell
+    torchrun --nnodes=1 --nproc_per_node=1 --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=localhost train.py \
+     --save-dir ... \
+     --model ... \
+     --start-lr ... \
+     --cook-config  ... \ # give your cook config path
+```
+
 
 <div id="quick-start"></div>
 
 ## Quick Start
 
-The `cpm_live_example` folder provides pruning example based on CPM-Live, please check [it](https://github.com/OpenBMB/BMCook/blob/main/cpm_live_example/README.md) for more details.
+The `examples` folder provides pruning example based on CPM-Live, GPT2-Base, T5-large, please check [examples](https://github.com/OpenBMB/BMCook/blob/main/examples/README.md) for more details.
 
-The `gpt-example` folder provides example codes based on GPT2-Base.
+Take GPT2 as example:
 
 Quantization-aware training：
 
@@ -131,15 +239,17 @@ Combine quantization, pruning and knowledge distillation：
 
 ## Performances
 
-Based on GPT-J, we evaluate different combinations of compression techniques. The corpus is OpenWebText. We also train a small GPT-J with 0.7B parameters based on this corpus from scratch, GPT-J (0.7B).
+Based on T5-3B, we evaluate different combinations of compression techniques. The corpus for compression is the Pile. The evaluation datasets includes SST-2, MNLI, and SQuAD. Specifically, we freeze the compressed models and adopt adapter-tuning.
 
-|                        |     LM Loss    |     Relative Performance    |     Speedup    |
+|                        |     Average Performance    |     Relative Performance    |     Speedup    |
 |------------------------|----------------|-----------------------------|----------------|
-|     GPT-J              |           3.37 |                        -    |          1x    |
-|     GPT-J (0.7B)       |           4.06 |                       83.0% |         ~10x   |
-|     GPT-J (P+D)        |           3.57 |                       94.4% |          2x    |
-|     GPT-J (P+D+Q)      |           3.58 |                       94.1% |          8x    |
-|     GPT-J (P+D+Q+M)    |           3.69 |                       91.3% |          10x   |
+|     T5-3B              |           0.9258 |                        -    |          1x    |
+|     T5-Base       |           0.8796 |                       95.0% |         7x   |
+|     T5-3B (P+D)        |           0.9150 |                       98.8% |          2x    |
+|     T5-3B (P+D+Q)      |           0.9126 |                       98.6% |          8x    |
+|     T5-3B (P+D+Q+M)    |           0.9017 |                       97.4% |          12x   |
+
+
 
 D denotes knowledge distillation. P denotes pruning. Q denotes quantization. M denotes MoEfication.
 
diff --git a/bmcook/__init__.py b/bmcook/__init__.py
new file mode 100644
index 0000000..3f52ad7
--- /dev/null
+++ b/bmcook/__init__.py
@@ -0,0 +1,10 @@
+from .trainer import CookTrainer, CPMAntTrainer
+from .store import save_spruned
+
+from .distilling import BMDistill
+from .pruning import BMPrune
+from .moe import BMMoE
+from .quant import BMQuant
+from .utils import arguments
+
+from .utils.config import ConfigParser
\ No newline at end of file
diff --git a/bmcook/distilling/__init__.py b/bmcook/distilling/__init__.py
index b5ebe49..662c2ef 100644
--- a/bmcook/distilling/__init__.py
+++ b/bmcook/distilling/__init__.py
@@ -1,11 +1,7 @@
 import types
-from numpy import record
-import torch
-from torch import nn
 import bmtrain as bmt
 import torch.nn.functional as F
-import cpm_kernels.torch as ct
-import model_center
+import model_center.layer as Layer
 
 
 class BMDistill:
@@ -14,7 +10,7 @@ class BMDistill:
     '''
 
     @classmethod
-    def set_forward(cls, student, teacher, foward_fn, config):
+    def set_forward(cls, student, teacher, foward_fn, config, target_linear = Layer.Linear):
         '''
         Modify the forward function of the student model to compute additional knowledge distillation loss.
 
@@ -37,60 +33,123 @@ def set_forward(cls, student, teacher, foward_fn, config):
             s_module_map, t_module_map = get_module_map(distill_config['mse_hidn_module'])
             update_forward(student, teacher, s_module_map, t_module_map)
 
-        def forward(model, enc_input, enc_length, dec_input, dec_length, targets, loss_func):
-
-            with bmt.inspect.inspect_tensor() as inspector:
-                outputs = foward_fn(
-                    model, enc_input, enc_length, dec_input, dec_length, targets, loss_func)
-                outputs_t = teacher(enc_input, enc_length, dec_input, dec_length, return_logits=True)
-
-            records = {}
-            for record in inspector._summary:
-                records[record['name']] = record['tensor']
-
-            loss = outputs[0]
-            model_outputs = outputs[1]
-            logits_s = model_outputs
-
-
-            # Compute loss and d_loss
-            d_loss = 0.0
-            if distill_config['ce_scale'] > 0:
-                temp = distill_config['ce_temp']
-                logits_t = outputs_t.detach()
-                prob_t = F.softmax(logits_t / temp, dim=-1)
-                log_prob_s = F.log_softmax(logits_s / temp, dim=-1)
-                d_loss += -(prob_t * log_prob_s).sum(dim=1).mean() * distill_config['ce_scale']
-        
-            # MSE loss 
-            if distill_config['mse_hidn_scale'] > 0:
-                for module_name in s_module_map:
-                    t_module_name = s_module_map[module_name]['t']['name']
-                    student_t = records[module_name+'_student']
-                    teacher_t = records[t_module_name+'_teacher'].detach()
-
-                    if distill_config['mse_hidn_proj']:
-                        if 'mapping' not in s_module_map[module_name]:
-                            t_dim = teacher_t.size(-1)
-                            s_dim = student_t.size(-1)
-                            # May be different on different devices
+        if cls.version:
+            def forward(model, loss_func, targets, *model_args, **model_kwargs):
+
+                with bmt.inspect.inspect_tensor() as inspector:
+                    outputs = foward_fn(
+                        model, loss_func, targets, *model_args, **model_kwargs)
+                    outputs_t = teacher(*model_args, **model_kwargs)
+
+                records = {}
+                for record in inspector._summary:
+                    records[record['name']] = record['tensor']
+
+                loss = outputs.loss
+                model_outputs = outputs.original_output
+                logits_s = model_outputs
+
+
+                # Compute loss and d_loss
+                d_loss = 0.0
+                if distill_config['ce_scale'] > 0:
+                    temp = distill_config['ce_temp']
+                    logits_t = outputs_t.detach()
+                    prob_t = F.softmax(logits_t / temp, dim=-1)
+                    log_prob_s = F.log_softmax(logits_s / temp, dim=-1)
+                    d_loss += -(prob_t * log_prob_s).sum(dim=1).mean() * distill_config['ce_scale']
+            
+                # MSE loss 
+                if distill_config['mse_hidn_scale'] > 0:
+                    for module_name in s_module_map:
+                        t_module_name = s_module_map[module_name]['t']['name']
+                        student_t = records[module_name+'_student']
+                        teacher_t = records[t_module_name+'_teacher'].detach()
+
+                        if distill_config['mse_hidn_proj']:
+                            if 'mapping' not in s_module_map[module_name]:
+                                t_dim = teacher_t.size(-1)
+                                s_dim = student_t.size(-1)
+                                # May be different on different devices
+                                
+                                s_module_map[module_name]['mapping'] = target_linear(t_dim, s_dim, init_std=0.02)
+                                bmt.init_parameters(s_module_map[module_name]['mapping'])
+                                s_module_map[module_name]['mapping'].to(teacher_t.device)
+                                bmt.synchronize()
                             
-                            s_module_map[module_name]['mapping'] = model_center.layer.Linear(t_dim, s_dim, init_std=0.02)
-                            bmt.init_parameters(s_module_map[module_name]['mapping'])
-                            s_module_map[module_name]['mapping'].to(teacher_t.device)
-                            bmt.synchronize()
-                        
-                        teacher_t = s_module_map[module_name]['mapping'](teacher_t)
+                            teacher_t = s_module_map[module_name]['mapping'](teacher_t)
+                            
+                        #normalize
+                        student_t_norm = student_t / (student_t.norm(dim=-1)).mean()
+                        teacher_t_norm = teacher_t / (teacher_t.norm(dim=-1)).mean()
                         
-                    cur_loss = (student_t - teacher_t).pow(2).mean() * distill_config['mse_hidn_scale']
-                    d_loss += cur_loss
+                        cur_loss = (student_t_norm - teacher_t_norm).pow(2).mean() * distill_config['mse_hidn_scale']
+                        d_loss += cur_loss
+                
+                loss = loss + d_loss
+
+                # update loss & append distillation loss
+                outputs.loss = loss
+                outputs.d_loss = d_loss
+                return outputs
+        else:
+            def forward(model, loss_func, targets, *model_args, **model_kwargs):
+
+                with bmt.inspect.inspect_tensor() as inspector:
+                    outputs = foward_fn(
+                        model, loss_func, targets, *model_args, **model_kwargs)    
+                    outputs_t = teacher(*model_args, **model_kwargs)
+
+                records = {}
+                for record in inspector._summary:
+                    records[record['name']] = record['tensor']
+
+                loss = outputs.loss
+                model_outputs = outputs.original_output
+                logits_s = model_outputs.logits
+
+                # Compute loss and d_loss
+                d_loss = 0.0
+                if distill_config['ce_scale'] > 0:
+                    temp = distill_config['ce_temp']
+                    logits_t = outputs_t.logits.detach()
+                    prob_t = F.softmax(logits_t / temp, dim=-1)
+                    log_prob_s = F.log_softmax(logits_s / temp, dim=-1)
+                    d_loss += -(prob_t * log_prob_s).sum(dim=1).mean() * distill_config['ce_scale']
             
-            loss = loss + d_loss
+                # MSE loss 
+                if distill_config['mse_hidn_scale'] > 0:
+                    for module_name in s_module_map:
+                        t_module_name = s_module_map[module_name]['t']['name']
+                        student_t = records[module_name+'_student']
+                        teacher_t = records[t_module_name+'_teacher'].detach()
+
+                        if distill_config['mse_hidn_proj']:
+                            if 'mapping' not in s_module_map[module_name]:
+                                t_dim = teacher_t.size(-1)
+                                s_dim = student_t.size(-1)
+                                # May be different on different devices
+                                
+                                s_module_map[module_name]['mapping'] = target_linear(t_dim, s_dim, init_std=0.02)
+                                bmt.init_parameters(s_module_map[module_name]['mapping'])
+                                s_module_map[module_name]['mapping'].to(teacher_t.device)
+                                bmt.synchronize()
+                            
+                            teacher_t = s_module_map[module_name]['mapping'](teacher_t)
 
-            # update loss & append distillation loss
-            outputs[0] = loss
-            outputs = outputs + [d_loss, ]
-            return outputs
+                        #normalize
+                        student_t_norm = student_t / (student_t.norm(dim=-1)).mean()
+                        teacher_t_norm = teacher_t / (teacher_t.norm(dim=-1)).mean()
+                        
+                        cur_loss = (student_t_norm - teacher_t_norm).pow(2).mean() * distill_config['mse_hidn_scale']
+                        d_loss += cur_loss
+                
+                loss = loss + d_loss
+
+                # update loss & append distillation loss
+                outputs.loss = loss
+                outputs.d_loss = d_loss
+                return outputs
         return forward
 
 def get_module_info(info):
diff --git a/bmcook/moe/__init__.py b/bmcook/moe/__init__.py
index bde453a..629cf1d 100644
--- a/bmcook/moe/__init__.py
+++ b/bmcook/moe/__init__.py
@@ -1,8 +1,4 @@
-import numpy as np
-import torch
-import cpm_kernels.torch as ct
 import types
-import os
 import bmtrain as bmt
 
 class BMMoE:
@@ -29,17 +25,18 @@ def get_hidden(model, config, forward_fn):
 
         update_forward(modules)
 
-        def forward(model, enc_input, enc_length, dec_input, dec_length, targets, loss_func):
+        def forward(model, loss_func, targets, *model_args, **model_kwargs):
             with bmt.inspect.inspect_tensor() as inspector:
                 outputs = forward_fn(
-                    model, enc_input, enc_length, dec_input, dec_length, targets, loss_func)
+                    model, loss_func, targets, *model_args, **model_kwargs)
             
             records = {}
             for record in inspector._summary:
                 if 'moe_hidden' in record['name']:
                     records[record['name']] = record['tensor']
             
-            return outputs + [records]
+            outputs.moe_records = records
+            return outputs
         return forward
 
     # @staticmethod
diff --git a/bmcook/pruning/__init__.py b/bmcook/pruning/__init__.py
index 4084604..aec751c 100644
--- a/bmcook/pruning/__init__.py
+++ b/bmcook/pruning/__init__.py
@@ -2,13 +2,12 @@
 import types
 import torch
 import bmtrain as bmt
+import os
+import json
 from bmtrain.block_layer import storage_type_cuda, round_up
 from .prune_func import m4n2_1d, m4n2_2d_greedy
-import os
-import sys
-sys.path.append('..')  # add path to make docs
-from bmcook.utils.config import ConfigParser
-from .sprune import L0_Module_coarse, L0_Module_fine
+from .sprune import SPruneEngine, SPrunePlugin
+
 
 def get_trivial_mask(p):
     return torch.ones_like(p)
@@ -52,13 +51,10 @@ def mask_storage(ordered_masks, storage_params, storage_info):
             to_offset_end = offset_end + param_st - storage_st
 
             # copy to buffer
-            #storaged_mask[storage_type].storage()[to_offset_st: to_offset_end].copy_(contiguous_param.storage()[offset_st: offset_end])
             d_dtype = storaged_mask[storage_type].dtype
             d_device = storaged_mask[storage_type].device
             contiguous_param = contiguous_param.to(device=d_device)
             
-            #if config['rank'] == 0:
-            #    print(config['rank'], d_device, contiguous_param.device)
             assert d_device == contiguous_param.device, "The devices does not match, which is not allowed when duplicating storage."
             torch.tensor([], dtype=d_dtype, device=d_device).set_(storaged_mask[storage_type].storage(), to_offset_st, (to_offset_end - to_offset_st,))[:] = \
                         torch.tensor([], dtype=d_dtype, device=d_device).set_(contiguous_param.storage(), offset_st, (offset_end - offset_st,))[:]
@@ -76,7 +72,7 @@ class BMPrune:
     _model = None
     _masks = None
     _optimizer = None
-    _sprune_module = None
+    _sprune = False
 
     @classmethod
     def compute_mask(cls, model, config):
@@ -103,26 +99,12 @@ def compute_mask(cls, model, config):
             func = m4n2_1d
         elif prune_config['mask_method'] == 'm4n2_2d':
             func = m4n2_2d_greedy
-        elif prune_config['mask_method'] == 'coarse-grained':
-            sprune_config = ConfigParser('config/l0_pruning.json').get('coarse-grained')
-            if not sprune_config['train_mask']:
-                mask = torch.load(sprune_config['coarse_mask'])
-                return False, mask
-            else:
-                sprune_module = L0_Module_coarse(model, sprune_config)
-                cls._optim_and_scheduler = sprune_module.create_sprune_optimizer()
-                cls._sprune_module = sprune_module
-                return True, None
-        elif prune_config['mask_method'] == 'fine-grained':
-            sprune_config = ConfigParser('config/l0_pruning.json').get('fine-grained')
-            if not sprune_config['train_mask']:
-                mask = torch.load(sprune_config['heads_mask'])
-                return False, mask
-            else:
-                sprune_module = L0_Module_fine(sprune_config)
-                cls._optim_and_scheduler = sprune_module.create_sprune_optimizer()
-                cls._sprune_module = sprune_module
-                return True, None
+        elif prune_config['mask_method'] == 'sprune':
+            sprune_config = prune_config['sprune']
+            plugin = SPrunePlugin(model)
+            cls.sprune_engine = SPruneEngine(sprune_config, plugin)
+            cls._sprune = True
+            return
         else:
             raise ValueError("Unknown mask method: {}".format(prune_config['mask_method']))
 
@@ -181,6 +163,29 @@ def compute_mask(cls, model, config):
 
         cls._masks = _masks
 
+    @classmethod
+    def set_forward_sprune(cls, forward_fn):
+        r"""
+        Modify the CookTrainer.forward
+
+        :param forward_fn: func CookTrainer.forward to modify.
+        :return: new forward modified
+        """
+        if cls._sprune is True:
+            def forward(model, loss_func, targets, *model_args, **model_kwargs):
+                outputs = forward_fn(model, loss_func, targets, *model_args, **model_kwargs)
+                loss = outputs.loss
+
+                lag_loss, sparsity = cls.sprune_engine.update()
+                # lag_loss should not be scaled
+                # loss += lag_loss
+                
+                outputs.loss, outputs.lag_loss, outputs.sparsity = loss, lag_loss, sparsity
+                return outputs
+        else:
+            forward = forward_fn
+        return forward
+
     @classmethod
     def set_optim_for_pruning(cls, optimizer):
         '''
@@ -198,27 +203,44 @@ def set_optim_for_pruning(cls, optimizer):
         cls._optimizer = optimizer
         cls._optimizer.step_old = optimizer.step
 
-        def _step(opt_self, *args, **kwargs):
-            # prune gradients before step method
-            with torch.no_grad():
-                for p, mask, flag in cls._masks:
-                    if flag:
-                        for k in p:
-                            if p[k].grad is not None:
-                                p[k].grad.mul_(mask[k])
-                    else:
-                        if p.grad is not None:
-                            p.grad.mul_(mask)
-            # call original optimizer step method
-            rval = opt_self.step_old(*args, **kwargs)
-            # prune parameters after step method
-            with torch.no_grad():
-                for p, mask, flag in cls._masks:
-                    if flag:
-                        for k in p:
-                            p[k].mul_(mask[k])
-                    else:
-                        tmp_mask = torch.tensor(mask, dtype=p.dtype, device=p.device)
-                        p.mul_(tmp_mask)
-            return rval
-        cls._optimizer.step = types.MethodType(_step, cls._optimizer)
+        if cls._sprune is False:
+            def _step(opt_self, *args, **kwargs):
+                # prune gradients before step method
+                with torch.no_grad():
+                    for p, mask, flag in cls._masks:
+                        if flag:
+                            for k in p:
+                                if p[k].grad is not None:
+                                    p[k].grad.mul_(mask[k])
+                        else:
+                            if p.grad is not None:
+                                p.grad.mul_(mask)
+                # call original optimizer step method
+                rval = opt_self.step_old(*args, **kwargs)
+                # prune parameters after step method
+                with torch.no_grad():
+                    for p, mask, flag in cls._masks:
+                        if flag:
+                            for k in p:
+                                p[k].mul_(mask[k])
+                        else:
+                            tmp_mask = torch.tensor(mask, dtype=p.dtype, device=p.device)
+                            p.mul_(tmp_mask)
+                return rval
+            cls._optimizer.step = types.MethodType(_step, cls._optimizer)
+        else:
+            cls._optimizer.zero_grad_old = optimizer.zero_grad
+
+            def _step(opt_self, *args, **kwargs):
+                rval = opt_self.step_old(*args, **kwargs)
+                cls.sprune_engine.sp_optimizer.step()
+                cls.sprune_engine.lagrangian_optimizer.step()
+                return rval
+            cls._optimizer.step = types.MethodType(_step, cls._optimizer)
+
+            def _zero_grad(opt_self, *args, **kwargs):
+                rval = opt_self.zero_grad_old(*args, **kwargs)
+                cls.sprune_engine.sp_optimizer.zero_grad()
+                cls.sprune_engine.lagrangian_optimizer.zero_grad()
+                return rval
+            cls._optimizer.zero_grad = types.MethodType(_zero_grad, cls._optimizer)
\ No newline at end of file
diff --git a/bmcook/pruning/l0_pruning.json b/bmcook/pruning/l0_pruning.json
deleted file mode 100644
index 50720a5..0000000
--- a/bmcook/pruning/l0_pruning.json
+++ /dev/null
@@ -1,19 +0,0 @@
-{
-    "coarse-grained": {
-        "train_mask": false,
-        "coarse_mask": "results/sprune/2B_layer_mask.pt",
-        "target_sparsity": 0.67
-    },
-
-    "fine-grained": {
-        "train_mask": true,
-        "coarse_mask": "results/sprune/2B_layer_mask.pt",
-        "num_heads": 16,
-        "dim_ff": 700,
-        "target_num_heads": 4,
-        "target_dimff": 350,
-        "heads_mask": "results/sprune/300M_head_mask.pt",
-        "dimff_mask": "results/sprune/300M_dimff_mask.pt",
-        "prune_mode": "att"
-    }
-}
\ No newline at end of file
diff --git a/bmcook/pruning/sprune.py b/bmcook/pruning/sprune.py
deleted file mode 100644
index 077fc37..0000000
--- a/bmcook/pruning/sprune.py
+++ /dev/null
@@ -1,248 +0,0 @@
-import torch
-import bmtrain as bmt
-import torch.nn as nn
-from torch.nn.parameter import Parameter
-
-limit_a, limit_b, epsilon = -.1, 1.1, 1e-6
-
-class L0_Module(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-        self.lambda_1 = Parameter(torch.tensor(0.0, dtype=torch.float, device='cuda'))
-        self.lambda_2 = Parameter(torch.tensor(0.0, dtype=torch.float, device='cuda'))
-
-    def _sample_z(self, loga, beta=2./3.):
-        eps = torch.FloatTensor(*loga.shape).uniform_(epsilon, 1-epsilon).to(loga.device)
-        s = torch.sigmoid((torch.log(eps / (1 - eps)) + loga) / beta)
-        s = s * (limit_b - limit_a) + limit_a
-        z = s.clamp(min=0., max=1.)
-        return z
-    
-    def cdf_qz(self, eps, loga, beta=2./3.):
-        """Implements the CDF of the 'stretched' concrete distribution"""
-        xn = (eps - limit_a) / (limit_b - limit_a)  # 0.083333
-        s = torch.sigmoid((torch.log(xn / (1 - xn)) * beta - loga))
-        z = s.clamp(min=epsilon, max=1 - epsilon)
-        return z
-
-    def initialize_parameters(self, shape, dtype=torch.float, init_mean=10., init_std=1e-2, requires_grad=True):
-        param = Parameter(torch.empty(shape, dtype=dtype), requires_grad=requires_grad)
-        param.data.normal_(init_mean, init_std)
-        param.requires_grad = True
-        
-        return param
-    
-    def create_sprune_optimizer(self):
-        l0_params = [{
-                    "params": [p for n, p in self.named_parameters() if "lambda" not in n],
-                    "weight_decay": 0.0,
-                    "lr": 0.1
-                }]
-        sp_optimizer = torch.optim.AdamW(l0_params)
-
-        lagrangian_params = [{
-                            "params": [p for n, p in self.named_parameters() if "lambda" in n],
-                            "weight_decay": 0.0,
-                            "lr": -0.1
-                            }]
-        lagrangian_optimizer = torch.optim.AdamW(lagrangian_params)
-
-        return (sp_optimizer, lagrangian_optimizer)
-
-
-class L0_Module_coarse(L0_Module):
-    '''
-    This L0_Module_coarse is for layer mask training
-    '''
-    def __init__(self, 
-                model : bmt.CheckpointBlock,
-                config
-                ) -> None:
-        super(L0_Module_coarse, self).__init__()
-        self.target_sparsity = config['target_sparsity']
-
-        self.masks = []
-        ordered_parameters = []
-        self.num_self_att_layer = 0
-        self.num_cross_att_layer = 0
-        self.num_ffn_layer = 0
-        self.prunable_all_params = 0
-        self.params_per_self_att = 0
-        self.params_per_cross_att = 0
-        self.params_per_ffn = 0
-
-        for prefix, block in model.named_modules():
-            if bmt.block_layer.CheckpointBlock == type(block):
-                block_for_info = block
-                ordered_parameters = [(k, v) for k, v in block.state_dict().items()]
-
-                self.prunable_all_params += sum([param['size'] for param in block._param_info if 'layernorm' not in param['name']])
-
-                for (k, v) in ordered_parameters:
-                    if 'self_att' in k:
-                        self.num_self_att_layer += 1
-                        break
-                for (k, v) in ordered_parameters:
-                    if 'ffn' in k:
-                        self.num_ffn_layer += 1
-                        break
-                if 'encoder' in prefix:
-                    pass
-                else:
-                    self.num_cross_att_layer += 1
-
-        for param in block_for_info._param_info:
-            if 'self_att' in param['name']:
-                self.params_per_self_att += param['size']
-            elif 'cross_att' in param['name']:
-                self.params_per_cross_att += param['size']
-            elif 'ffn' in param['name']:
-                self.params_per_ffn += param['size']
-        
-        self.self_att_layer_loga = self.initialize_parameters((self.num_self_att_layer))
-        self.cross_att_layer_loga = self.initialize_parameters((self.num_cross_att_layer))
-        self.ffn_layer_loga = self.initialize_parameters((self.num_ffn_layer))
-
-    def _deterministic_z(self, loga, size):
-        expected_num_nonzeros = torch.sum(1 - self.cdf_qz(torch.tensor(0.), loga))
-        expected_num_zeros = size - expected_num_nonzeros.item()
-        num_zeros = round(expected_num_zeros)
-        soft_mask = torch.ones_like(loga)
-        if num_zeros > 0:
-            if soft_mask.ndim == 0:
-                soft_mask = torch.tensor(0).to(loga.device)
-            else:
-                _, indices = torch.topk(loga, k=num_zeros, largest=False)  # 返回values, indices
-                soft_mask[indices] = 0.  # 置零
-        return soft_mask
-
-    def forward(self, training=True):
-        zs = {"self_att_layer_z": [], 
-                "cross_att_layer_z": [],
-                "ffn_layer_z": []}
-
-        if training:
-            bmt.print_rank(self.ffn_layer_loga)
-            zs["self_att_layer_z"] = self._sample_z(self.self_att_layer_loga).reshape((self.num_self_att_layer))
-            zs["cross_att_layer_z"] = self._sample_z(self.cross_att_layer_loga).reshape((self.num_cross_att_layer))
-            zs["ffn_layer_z"] = self._sample_z(self.ffn_layer_loga).reshape((self.num_ffn_layer))
-        else:
-            zs["self_att_layer_z"] = self._deterministic_z(self.self_att_layer_loga, self.num_self_att_layer).reshape((self.num_self_att_layer))
-            zs["cross_att_layer_z"] = self._deterministic_z(self.cross_att_layer_loga, self.num_cross_att_layer).reshape((self.num_cross_att_layer))
-            zs["ffn_layer_z"] = self._deterministic_z(self.ffn_layer_loga, self.num_ffn_layer).reshape((self.num_ffn_layer))
-        return zs
-
-    def get_num_parameters(self):
-        num_parameters = 0
-
-        self_att_score = 1 - self.cdf_qz(torch.tensor(0.), self.self_att_layer_loga)
-        cross_att_score = 1 - self.cdf_qz(torch.tensor(0.), self.cross_att_layer_loga)
-        ffn_score = 1 - self.cdf_qz(torch.tensor(0.), self.ffn_layer_loga)
-        
-        num_parameters += torch.sum(self_att_score) * self.params_per_self_att
-        num_parameters += torch.sum(cross_att_score) * self.params_per_cross_att
-        num_parameters += torch.sum(ffn_score) * self.params_per_ffn
-
-        return num_parameters
-
-    def lagrangian_regularization(self):
-        target_sparsity = self.target_sparsity  # 指定目标sparsity
-        expected_size = self.get_num_parameters()
-        
-        expected_sparsity = 1 - expected_size / self.prunable_all_params
-
-        lagrangian_loss = self.lambda_1 * (expected_sparsity - target_sparsity) + \
-                            self.lambda_2 * (expected_sparsity - target_sparsity) ** 2
-
-        return lagrangian_loss, expected_sparsity, target_sparsity
-
-
-class L0_Module_fine(L0_Module):
-    '''
-    This L0_Module_fine is for heads/dimff mask training given fixed layer mask
-    '''
-    def __init__(self, 
-                prune_config,
-                ):
-        super(L0_Module_fine, self).__init__()
-        self.layer_zs = torch.load(prune_config['coarse_mask'])
-        self.num_heads = prune_config['num_heads']
-        self.dim_ff = prune_config['dim_ff']
-        self.prune_mode = prune_config['prune_mode']
-        
-        self.att_zs = self.layer_zs['self_att_layer_z']
-        self.ffn_zs = self.layer_zs['ffn_layer_z']
-        self.num_att_layers = int(torch.sum(self.att_zs).item())
-        self.num_ffn_layers = int(torch.sum(self.ffn_zs).item())
-
-        if self.prune_mode == 'att':
-            self.heads_loga = self.initialize_parameters((self.num_att_layers, self.num_heads))
-            self.target_num_heads = prune_config['target_num_heads']
-        elif self.prune_mode == 'ffn':
-            mask_z = torch.load(prune_config['heads_mask'])
-            heads_z = mask_z['heads_z']  # (25, 32)
-            self.heads_loga = Parameter(heads_z, requires_grad=False)
-            self.dimff_loga = self.initialize_parameters((self.num_ffn_layers, self.dim_ff))
-            self.target_dimff = prune_config['target_dimff']
-
-    def _deterministic_z(self, loga, size):
-        expected_num_nonzeros = torch.sum(1 - self.cdf_qz(torch.tensor(0.), loga), -1)
-        expected_num_zeros = size - expected_num_nonzeros
-        res, num_zeros = [], []
-        if self.prune_mode == 'att':
-            num_zeros = self.num_heads - self.target_num_heads
-        elif self.prune_mode == 'ffn':
-            num_zeros = self.dim_ff - self.target_dimff
-        for index in range(len(expected_num_zeros)):
-            cur_layer = loga[index]
-            soft_mask = torch.ones_like(cur_layer)
-            if num_zeros > 0:
-                if soft_mask.ndim == 0:
-                    soft_mask = torch.tensor(0).to(cur_layer.device)
-                else:
-                    _, indices = torch.topk(cur_layer, k=num_zeros, largest=False)  # 返回values, indices
-                    soft_mask[indices] = 0.  # 置零
-            res.append(soft_mask)
-        return torch.stack(res)
-
-    def forward(self, training=True):
-        if self.prune_mode == 'att':
-            zs = {"heads_z" : self.heads_loga}
-            if training:
-                zs["heads_z"] = self._sample_z(self.heads_loga).reshape((self.num_att_layers, self.num_heads))
-            else:
-                zs["heads_z"] = self._deterministic_z(self.heads_loga, 25).reshape((self.num_att_layers, self.num_heads))
-
-        elif self.prune_mode == 'ffn':
-            zs = {"heads_z" : self.heads_loga, "dimff_z": self.dimff_loga,}
-            if training:
-                zs["dimff_z"] = self._sample_z(self.dimff_loga).reshape((self.num_ffn_layers, self.dim_ff))
-            else:
-                zs["dimff_z"] = self._deterministic_z(self.dimff_loga, self.num_ffn_layers).reshape((self.num_ffn_layers, self.dim_ff))
-        
-        return zs
-
-    def lagrangian_regularization(self):
-        if self.prune_mode == 'att':
-            target_num_heads = self.target_num_heads
-            heads_score = 1 - self.cdf_qz(torch.tensor(0.), self.heads_loga).cuda()
-            expected_num_heads = torch.sum(heads_score, -1)
-            loss_heads = torch.sum((target_num_heads - expected_num_heads) / self.num_heads)
-
-            lagrangian_loss = self.lambda_1 * (loss_heads) + \
-                                self.lambda_2 * (loss_heads) ** 2
-            expected_num = expected_num_heads
-            target = target_num_heads
-        
-        elif self.prune_mode == 'ffn':
-            target_dimff = self.target_dimff
-            dimff_score = 1 - self.cdf_qz(torch.tensor(0.), self.dimff_loga).cuda()
-            expected_num_dimff = torch.sum(dimff_score, -1)
-            loss_dimff = torch.sum((target_dimff - expected_num_dimff) / self.dim_ff)
-            lagrangian_loss = self.lambda_1 * (loss_dimff) + \
-                                self.lambda_2 * (loss_dimff) ** 2
-            expected_num = expected_num_dimff
-            target = target_dimff
-
-        return lagrangian_loss, expected_num, target
-
diff --git a/bmcook/pruning/sprune/README.md b/bmcook/pruning/sprune/README.md
new file mode 100644
index 0000000..514bfb5
--- /dev/null
+++ b/bmcook/pruning/sprune/README.md
@@ -0,0 +1,31 @@
+# BMCook SPrune
+
+## Overview
+Structure pruning toolkit of BMCook.
+
+If you want to prune a large PLM in a structured way, BMCook SPrune may help to get a desirable subtructure, based on L0 regularization and lagrangian method. For details, refer to:
+
+1. [Learning Sparse Neural Networks through L_0 Regularization](https://openreview.net/forum?id=H1Y8hhg0b)
+
+2. [Structured Pruning of Large Language Models](https://arxiv.org/abs/1910.04732)
+
+BMCook SPrune use two classes to manage structure pruning process: **SPrunePlugin** and **SPruneEngine**.
+
+**SPrunePlugin** consists of the mask, and directly related to the model. It's responsible for pruning. **SPruneEngine** manage the mask with specific strategy and criterion. It's responsible for finding an appropriate mask.
+
+## Usage
+
+1. Initialize a plugin from the target compressing model:
+```python
+plugin = SPrunePlugin(model)
+```
+2. Initialize an engine from the plugin and config:
+```python
+engine = SPruneEngine(sprune_config, plugin)
+```
+
+The process is coupled with BMCook.pruning. You can find it in "pruning/\_\_init\_\_.py"
+
+## Future
+
+We will provide more sprune strategies, like iterational pruning and hidden pruning, in the future. Stay tuned！ 
\ No newline at end of file
diff --git a/bmcook/pruning/sprune/__init__.py b/bmcook/pruning/sprune/__init__.py
new file mode 100644
index 0000000..21fd318
--- /dev/null
+++ b/bmcook/pruning/sprune/__init__.py
@@ -0,0 +1,2 @@
+from .plugin import SPrunePlugin
+from .engine import SPruneEngine
\ No newline at end of file
diff --git a/cpm_live_example/do_prune.py b/bmcook/pruning/sprune/do_prune.py
similarity index 100%
rename from cpm_live_example/do_prune.py
rename to bmcook/pruning/sprune/do_prune.py
diff --git a/bmcook/pruning/sprune/engine.py b/bmcook/pruning/sprune/engine.py
new file mode 100644
index 0000000..734ae20
--- /dev/null
+++ b/bmcook/pruning/sprune/engine.py
@@ -0,0 +1,188 @@
+import torch
+import bmtrain as bmt
+from typing import Dict
+from torch.nn.parameter import Parameter
+from .func import determinate_mask, sample, binarize
+from .plugin import SPrunePlugin
+from .utils import get_params_from_block
+
+
+class SPruneStrategy:
+    def __init__(self, config: Dict) -> None:
+        self.criterion = config['criterion']
+        assert self.criterion == 'l0', "BMCook sprune do not support other criterions besides l0 yet."
+        self.fixed_mask_path = config['fixed_mask_path']
+        self.training_mask = config['training_mask']
+        self.mask_mode = config['mask_mode']
+        self.target_mode = config['target_mode']
+        self.target_sparsity = config['target_sparsity']
+
+
+class SPruneEngine:
+    r"""
+    SPruneEngine is used for the mask computation and update of SPrunePlugin.
+
+    The engine design is based on L0 regularization method and a lagrangian term. For L0 regularization details, see paper
+        "Learning Sparse Neural Networks through L_0 Regularization" <https://openreview.net/forum?id=H1Y8hhg0b>.
+        For lagrangian term in PLM structure pruning, see paper "Structured Pruning of Large Language Models" 
+        <https://arxiv.org/abs/1910.04732>.
+    """
+    def __init__(self, config: Dict, plugin: SPrunePlugin) -> None:
+        r"""Init the SpruneEngine from a SPrunePlugin. It will initilize all the :class:`torch.nn.Parameter`
+        used for learning the sprune mask, and create the optimizer for l0 regularization.
+
+        Args:
+            config: `(Dict)`, the sprune config.
+            plugin: `(SPrunePlugin)`, the SPrunePlugin.
+        """
+        super().__init__()
+        self.strategy = SPruneStrategy(config)
+        self.target_sparsity = self.strategy.target_sparsity
+        self.plugin = plugin
+        self.training = True
+
+        self.lambda_1 = Parameter(torch.tensor(0., dtype=torch.float, device='cuda'))
+        self.lambda_2 = Parameter(torch.tensor(0., dtype=torch.float, device='cuda'))
+        self.training_loga = {}
+        for mask in self.strategy.training_mask:
+            shape = self.plugin.info_to_engine['shape'][mask]
+            self.training_loga[mask+'_loga'] = Parameter(torch.empty(shape[0], dtype=torch.float, device='cuda').normal_(0., 1e-2))
+
+        self.create_sprune_optimizer()
+
+    def create_sprune_optimizer(self):
+        r"""Create the sprune optimizer and lagrangian optimizer, making the learning of loga and 
+        lagrangian terms to be an adversarial game.
+        
+        sprune optimizer will manage the loga parameters.
+
+        lagrangian optimizer will manage the lagrangian terms.
+        """
+        l0_params = [{
+                        "params": [p for _, p in self.training_loga.items()],
+                        "weight_decay": 0.0,
+                        "lr": 0.1
+                        }]
+        self.sp_optimizer = torch.optim.AdamW(l0_params)
+
+        lagrangian_params = [{
+                    "params": [self.lambda_1, self.lambda_2],
+                    "weight_decay": 0.0,
+                    "lr": -0.1
+                }]
+        self.lagrangian_optimizer = torch.optim.AdamW(lagrangian_params)
+    
+    def update(self):
+        r"""
+        update the sprune parameters and lagrangian parameters.
+        """
+        if self.training:
+            info_list = self.update_plugin_mask(training=True)
+            loss, sparsity = self.loss(info_list)
+            if torch.abs(sparsity - self.target_sparsity) < 5e-5:
+                bmt.print_rank("binarize the mask and begin finetune...")
+                info_list = self.update_plugin_mask(training=False)
+                self.lambda_1.requires_grad_(False)
+                self.lambda_2.requires_grad_(False)
+                for v in self.training_loga.values():
+                    v.requires_grad_(False)
+                self.training = False
+        else:
+            info_list = self.update_plugin_mask(training=False)
+            loss, sparsity = self.loss(info_list)
+        return loss, sparsity
+
+    def step(self):
+        r"""run :method:`.step()` of sprune optimizer and lagrangian optimizer"""
+        self.sp_optimizer.step()
+        self.lagrangian_optimizer.step()
+    
+    def zero_grad(self):
+        r"""run :method:`.zero_grad()` of sprune optimizer and lagrangian optimizer"""
+        self.sp_optimizer.zero_grad()
+        self.lagrangian_optimizer.zero_grad()
+
+    def loss(self, info_list):
+        r"""calculate the lagrangian loss. It can be calculated in sparsity(`float`) or dimension(`int`)"""
+        if self.strategy.target_mode == 'sparsity':
+            return self.lagrangian_loss_sparsity(info_list, layer_constraint=False)
+        elif self.strategy.target_mode == 'dimension':
+            return self.lagrangian_loss_dimension()
+
+    def update_plugin_mask(self, training: bool = True):
+        r"""update the mask managed in plugin"""
+        info_list = {}
+        for k, v in self.training_loga.items():
+            module = k.split('_loga')[0]
+
+            mask = sample(v) if training is True else binarize(v)
+            train_mask = determinate_mask(v)
+            assert mask.size(0) == train_mask.size(0)
+            
+            for index in range(mask.size(0)):
+                self.plugin.__dict__[module][index]['mask'] = mask[index].clone().detach()
+                
+                param = self.plugin.__dict__[module][index]['param']
+                index_all = self.plugin.__dict__[module][index]['index']
+
+                if index_all not in info_list:
+                    info_list[index_all] = {'module': [module], 'param': [param], 'score': [train_mask[index]]}
+                else:
+                    if module in info_list[index_all]['module']:
+                        module_correct = 'cross_' + module
+                    else:
+                        module_correct = module
+                    info_list[index_all]['module'].append(module_correct)
+                    info_list[index_all]['param'].append(param)
+                    info_list[index_all]['score'].append(train_mask[index])
+
+        return info_list
+
+    def lagrangian_loss_sparsity(self, info_list, layer_constraint: bool = False):
+        r"""The func 'lagrangian_loss_sparsity' is to calculate the lagrangian loss to get the target sparsity"""
+        expected_sparsity = get_params_from_block(info_list)
+        loss_sparsity = expected_sparsity - self.target_sparsity
+        if layer_constraint:
+            loss_sparsity = torch.mean(torch.abs(loss_sparsity))
+            expected_sparsity = torch.mean(expected_sparsity)
+
+        lagrangian_loss = self.lambda_1 * loss_sparsity + self.lambda_2 * (loss_sparsity ** 2)
+
+        return lagrangian_loss, expected_sparsity
+
+    def lagrangian_loss_dimension(self):
+        r"""calculate the lagrangian loss to get the target dimension"""
+        dimension_score = determinate_mask(self.training_loga)
+        all_dimension = dimension_score.size(1)
+        
+        expected_dimension = torch.sum(dimension_score, -1)
+        loss_dimension = torch.sum((self.target_dimension - expected_dimension) / all_dimension)
+        
+        lagrangian_loss = self.lambda_1 * loss_dimension + self.lambda_2 * (loss_dimension ** 2)
+        
+        return lagrangian_loss, expected_dimension
+
+    def get_model_sparsity(self, layer_constraint: bool = False):
+        r"""calculate the current sparsity to calculate lagrangian loss."""
+        if not layer_constraint:
+            total_res = 0
+            num_res = 0
+            for k, v in self.training_loga.items():
+                # layer-wise
+                for layer_index in range(v.size(0)):
+                    param = self.plugin.__dict__[k.split('_loga')[0]][layer_index]['param']
+                    v_cur = v[layer_index]
+
+                    total_res +=  torch.sum(determinate_mask(v_cur)) * param * 3
+                    num_res   +=  v_cur.numel() * param * 3
+            ratio = total_res / num_res
+        else:
+            ratio = []
+            for k, v in self.training_loga.items():
+                # layer-wise
+                for layer_index in range(v.size(0)):
+                    param = self.plugin.__dict__[k.split('_loga')[0]][layer_index]['param']
+                    v_cur = v[layer_index]
+                    ratio.append(torch.sum(determinate_mask(v_cur)) / v_cur.numel())
+                ratio = torch.stack(ratio)
+        return 1 - ratio  # return sparsity
diff --git a/bmcook/pruning/sprune/func.py b/bmcook/pruning/sprune/func.py
new file mode 100644
index 0000000..2d169fb
--- /dev/null
+++ b/bmcook/pruning/sprune/func.py
@@ -0,0 +1,91 @@
+import torch
+import math
+from torch import Tensor
+
+gamma, zeta, epsilon = -.1, 1.1, 1e-6
+beta = 2./3.
+
+def cdf_concrete_dist(eps: Tensor, loga: Tensor):
+    r"""Implements the CDF of the 'stretched' concrete distribution"""
+    xn = (eps - gamma) / (zeta - gamma)
+    s = torch.sigmoid((torch.log(xn / (1 - xn)) * beta - loga))
+    z = s.clamp(min=0., max=1.)
+    return z
+
+
+def hard_concrete_dist(loga: Tensor):
+    r"""Implements the gard concrete distribution. For details, see paper 
+    'Learning Sparse Neural Networks through L_0 Regularization' <https://openreview.net/forum?id=H1Y8hhg0b>."""
+    u = torch.FloatTensor(*loga.shape).uniform_(epsilon, 1-epsilon).to(loga.device)
+    s = torch.sigmoid((torch.log(u / (1 - u)) + loga) / beta)
+    s_shift = s * (zeta - gamma) + gamma
+    z = s_shift.clamp(min=0., max=1.).to(device='cuda', dtype=torch.half) 
+    return z
+
+def l0_norm_term(loga: Tensor):
+    r"""calculate the l0 norm term. For details, see paper
+    'Structured Pruning of Large Language Models' <https://arxiv.org/abs/1910.04732>"""
+    x = loga - beta * (math.log(- gamma / zeta))
+    loss = torch.sigmoid(x).sum()
+    return loss
+
+def determinate_mask(loga: Tensor):
+    sig = torch.sigmoid(loga)
+    s = sig * (zeta - gamma) + gamma
+    out = s.clamp(min=0., max=1.)
+    return out
+
+def sample(loga: Tensor):
+    eps = torch.FloatTensor(*loga.shape).uniform_(epsilon, 1-epsilon).to(loga.device)
+    s = torch.sigmoid((torch.log(eps / (1 - eps)) + loga) / beta)
+    s = s * (zeta - gamma) + gamma
+    z = s.clamp(min=0., max=1.)
+    z_ = z.to(device='cuda', dtype=torch.half)  # remove from the computation graph of sp_module
+    return z_
+
+def binarize(loga: Tensor):
+    mask = determinate_mask(loga)
+    expected_num_nonzeros = torch.sum(mask, -1)
+    total_num_nonzeros = loga.size(-1)
+    expected_num_zeros = total_num_nonzeros - expected_num_nonzeros
+    if loga.dim() == 1:
+        num_zeros = round(expected_num_zeros.item())
+        soft_mask = torch.ones_like(loga, dtype=torch.half)
+        if num_zeros > 0:
+            if soft_mask.ndim == 0:
+                soft_mask = torch.tensor(0).to(loga.device)
+            else:
+                _, indices = torch.topk(loga, k=num_zeros, largest=False)  # return values, indices
+                soft_mask[indices] = 0.  # set zero
+        res = soft_mask.to(device='cuda')  # remove from the computation graph of sp_module
+    else:
+        res = []
+        for index in range(loga.size(0)):
+            num_zero = round(expected_num_zeros[index].item())
+            cur_layer = loga[index]
+            soft_mask = torch.ones_like(cur_layer, dtype=torch.half)
+            if num_zero > 0:
+                if soft_mask.ndim == 0:
+                    soft_mask = torch.tensor(0).to(cur_layer.device)
+                else:
+                    _, indices = torch.topk(mask[index], k=num_zero, largest=False)  # 返回values, indices
+                    soft_mask[indices] = 0.  # 置零
+            res.append(soft_mask)
+        res = torch.stack(res)
+    return res
+
+
+def binarize_mask_1d(loga: Tensor):
+    expected_num_nonzeros = torch.sum(1 - cdf_concrete_dist(torch.tensor(0.), loga), -1)
+    expected_num_zeros = loga.size(1) - expected_num_nonzeros
+    
+    num_zeros = round(expected_num_zeros)
+    soft_mask = torch.ones_like(loga)
+    if num_zeros > 0:
+        if soft_mask.ndim == 0:
+            soft_mask = torch.tensor(0).to(loga.device)
+        else:
+            _, indices = torch.topk(loga, k=num_zeros, largest=False)  # return values, indices
+            soft_mask[indices] = 0.  # set zero
+    soft_mask = soft_mask.to(device='cuda', dtype=torch.half)  # remove from the computation graph of sp_module
+    return soft_mask
\ No newline at end of file
diff --git a/bmcook/pruning/sprune/plugin.py b/bmcook/pruning/sprune/plugin.py
new file mode 100644
index 0000000..2af969b
--- /dev/null
+++ b/bmcook/pruning/sprune/plugin.py
@@ -0,0 +1,261 @@
+import torch
+import bmtrain as bmt
+from model_center.layer import TransformerBlock
+from model_center.model import BaseModel
+
+from .utils import *
+
+class SPrunePlugin:
+    r"""
+    SPrunePlugin is a base class for structure prune in BMCook.
+    All the modules supported by SprunePlugin includes: TransformerBlock layer, Attention layer, Feedforward layer, num_heads, dim_head, dim_ff.
+    """
+    def __init__(self, model: BaseModel):
+        '''
+        analyze the structure prune methed.
+        '''
+        # model analysis
+        num_encoder_layers, num_decoder_layers = 0, 0
+        prunable_all_params, self_att_num = 0, 0
+        TRANSFORMER_MASK, FFN_MASK, ATT_MASK = [], [], []
+        NUM_HEADS_MASK, DIM_HEAD_MASK, DIM_FF_MASK = [], [], []
+        cross_att_buffer, cross_num_heads_buffer, cross_dim_head_buffer, names_buffer = [], [], [], []
+        for name, module in model.named_modules():
+            if type(module) in (bmt.block_layer.CheckpointBlock, TransformerBlock):
+                block_type, overall_index = name.split('.')[0], int(name.split('.')[2])
+                if block_type == 'encoder':
+                    num_encoder_layers += 1
+                if block_type == 'decoder':
+                    num_decoder_layers += 1
+
+                ordered_parameters = [(k, v) for k, v in module.state_dict().items()]
+                self_att_param, cross_att_param, ffn_param = 0, 0, 0
+                for (k, v) in ordered_parameters:
+                    if 'self_att' in k:
+                        self_att_param += v.numel()
+                    elif 'cross_att' in k:
+                        cross_att_param += v.numel()
+                    elif 'ffn' in k:
+                        ffn_param += v.numel()
+
+                dim_head = module.self_att.self_attention.dim_head
+                num_heads = module.self_att.self_attention.num_heads
+                dim_model = module.self_att.self_attention.dim_model
+                dim_ff = get_dim_ff(module.ffn.ffn)
+                
+                # model ststistics
+                transformer_layer_param = self_att_param + cross_att_param + ffn_param
+                prunable_all_params += transformer_layer_param
+
+                if transformer_layer_param > 0:
+                    TRANSFORMER_MASK.append({
+                                            'index': overall_index,
+                                            'param': transformer_layer_param,
+                                            'dim': 1,
+                                            'mask': None
+                                            })
+                    set_pruning_transformer(module, len(TRANSFORMER_MASK)-1, TRANSFORMER_MASK, is_bmtCBlock=False)
+
+                    if self_att_param > 0:
+                        ATT_MASK.append({
+                                                'index': overall_index, 
+                                                'param': self_att_param,
+                                                'dim': 1,
+                                                'mask': None,
+                                                })
+                        NUM_HEADS_MASK.append({
+                                                'index': overall_index, 
+                                                'param': dim_head,
+                                                'dim': num_heads,
+                                                'mask': None
+                                                })
+                        DIM_HEAD_MASK.append({
+                                                'index': overall_index, 
+                                                'param': num_heads,
+                                                'dim': dim_head,
+                                                'mask': None
+                                                })
+                        set_pruning_att(module.self_att, len(ATT_MASK)-1, ATT_MASK, NUM_HEADS_MASK, DIM_HEAD_MASK)
+                        self_att_num += 1
+                                           
+                    if cross_att_param > 0:
+                        names_buffer.append(name)
+                        cross_att_buffer.append({
+                                                'index': overall_index, 
+                                                'param': cross_att_param,
+                                                'dim': 1,
+                                                'mask': None,
+                                                })
+                        cross_num_heads_buffer.append({
+                                                    'index': overall_index, 
+                                                    'param': dim_head,
+                                                    'dim': num_heads,
+                                                    'mask': None
+                                                    })
+                        cross_dim_head_buffer.append({
+                                                    'index': overall_index, 
+                                                    'param': num_heads,
+                                                    'dim': dim_head,
+                                                    'mask': None
+                                                    })
+
+                    if ffn_param > 0:
+                        FFN_MASK.append({
+                                        'index': overall_index, 
+                                        'param': ffn_param,
+                                        'dim': 1,
+                                        'mask': None,
+                                        })
+                        DIM_FF_MASK.append({
+                                            'index': overall_index, 
+                                            'param': dim_model,
+                                            'dim': dim_ff,
+                                            'mask': None
+                                            })
+                        set_pruning_ffn(module.ffn, len(FFN_MASK)-1, FFN_MASK, DIM_FF_MASK)
+
+        # append cross_att to att and set pruning
+        for (module_name, cross_att, cross_num_heads, cross_dim_head) in \
+            zip(names_buffer, cross_att_buffer, cross_num_heads_buffer, cross_dim_head_buffer):
+
+            ATT_MASK.append(cross_att)
+            NUM_HEADS_MASK.append(cross_num_heads)
+            DIM_HEAD_MASK.append(cross_dim_head)
+            set_pruning_att(model.get_submodule(module_name).cross_att, len(ATT_MASK)-1, ATT_MASK, NUM_HEADS_MASK, DIM_HEAD_MASK)
+
+        # check exception
+        if TRANSFORMER_MASK == []:
+            raise TypeError("plugin doesn't maintain any mask, all the mask lists are empty, \
+                            please check if your model has the module type: bmt.CheckpointBlock or model_center.layer.TransformerBlock")
+        elif any((FFN_MASK == [], ATT_MASK == [], NUM_HEADS_MASK ==[], DIM_HEAD_MASK == [], DIM_FF_MASK == [])):
+            raise ValueError("Now BMCook doesn't support to prune model without feedforward layer or attention layer. It's also not allowed only layernorm parameters exist in these layers.")
+        
+        # init mask shape for the use of loga
+        transformer_mask_shape = (len(TRANSFORMER_MASK))
+        
+        num_heads_list = [mask['dim'] for mask in NUM_HEADS_MASK]
+        num_heads_layers, max_num_heads = len(NUM_HEADS_MASK), max(num_heads_list)
+        num_heads_shape = (num_heads_layers, max_num_heads)
+        num_heads_shape_mask = torch.stack(
+            [(torch.arange(max_num_heads) < att['dim']).long() \
+                for att in NUM_HEADS_MASK]
+            )
+        
+        dim_head_list = [mask['dim'] for mask in DIM_HEAD_MASK]
+        dim_head_layers, max_dim_head = len(DIM_HEAD_MASK), max(dim_head_list)
+        dim_head_shape = (dim_head_layers, max_dim_head)
+        dim_head_shape_mask = torch.stack(
+            [(torch.arange(max_dim_head) < att['dim']).long() \
+                for att in DIM_HEAD_MASK]
+        )
+        
+        ffn_list = [mask['dim'] for mask in DIM_FF_MASK]
+        ffn_layers, max_dim_ff = len(DIM_FF_MASK), max(ffn_list)
+        ffn_shape = (ffn_layers, max_dim_ff)
+        ffn_shape_mask = torch.stack(
+            [(torch.arange(max_dim_ff) < ffn['dim']).long() \
+                for ffn in FFN_MASK]
+        )
+
+        self.info_to_engine = {
+            'all_params': prunable_all_params,
+            'self_att_num': self_att_num,
+            'num_encoder_layers': num_encoder_layers,
+            'shape':{
+                'transformer': ((transformer_mask_shape), torch.ones(transformer_mask_shape)),
+                'att': ((num_heads_layers), torch.ones(num_heads_layers)),
+                'ffn': ((ffn_layers), torch.ones(ffn_layers)),
+                'num_heads': (num_heads_shape, num_heads_shape_mask),
+                'dim_head': (dim_head_shape, dim_head_shape_mask),
+                'dim_ff': (ffn_shape, ffn_shape_mask)
+            }
+        }
+
+        self.transformer = TRANSFORMER_MASK
+        self.att = ATT_MASK
+        self.ffn = FFN_MASK
+        self.num_heads = NUM_HEADS_MASK
+        self.dim_head = DIM_HEAD_MASK
+        self.dim_ff = DIM_FF_MASK
+
+        del num_heads_list, dim_head_list, ffn_list, names_buffer, cross_att_buffer, cross_num_heads_buffer, cross_dim_head_buffer
+
+    def print_masks(self, key: str = None):
+        r"""print the masks managed in SPrunePlugin"""
+        res = {'transformer': self.transformer,
+                'att': self.att,
+                'ffn': self.ffn,
+                'num_heads': self.num_heads,
+                'dim_head': self.dim_head,
+                'dim_ff': self.dim_ff}
+        if key is None:
+            print(res)
+        else:
+            print(res[key])
+
+    def save_plugin(self, path):
+        r"""save the plugin as a dict.
+        Args:
+            path: `(str)`, the save path.
+        """
+        res = {'transformer': self.transformer,
+                'att': self.att,
+                'ffn': self.ffn,
+                'num_heads': self.num_heads,
+                'dim_head': self.dim_head,
+                'dim_ff': self.dim_ff,
+                'info_to_engine': self.info_to_engine}
+        torch.save(res, path)
+
+    def load_plugin(self, path):
+        r"""load the saved dict to this plugin.
+        Args:
+            path: `(str)`, the file path.
+        """
+        plugin = torch.load(path)
+        for k, v in plugin.items():
+            self.__dict__[k] = v
+
+    def training_masks(self):
+        r"""
+        traverse all the necessary masks in this training.
+        """
+        for name in self.info_to_engine['shape'].keys():
+            for i, v in enumerate(self.__dict__[name]):
+                if v['mask'] is not None:
+                    yield name + '.' + str(i)
+    
+    def training_modules(self):
+        r"""
+        traverse all the necessary modules in this training.
+        """
+        for name in self.info_to_engine['shape'].keys():
+            for _, v in enumerate(self.__dict__[name]):
+                if v['mask'] is not None:
+                    yield name
+                    break
+
+    def get_sparsity(self):
+        r"""
+        calculate the sparsity in single grain
+        """
+        info_list = {}
+
+        for name in self.training_masks():
+            module, index = name.split('.')[0], int(name.split('.')[1])
+            param = self.__dict__[module][index]['param']
+            mask = self.__dict__[module][index]['mask']
+            index = self.__dict__[module][index]['index']
+            if index not in info_list:
+                info_list[index] = {'module': [module], 'param': [param], 'score': [mask]}
+            else:
+                if module in info_list[index]['module']:
+                    module = 'cross_' + module
+                info_list[index]['module'].append(module)
+                info_list[index]['param'].append(param)
+                info_list[index]['score'].append(mask)
+        
+        expected_params, all_params = get_params_from_block(info_list)
+
+        sparsity = 1 - expected_params / all_params
+        return sparsity
\ No newline at end of file
diff --git a/bmcook/pruning/sprune/utils.py b/bmcook/pruning/sprune/utils.py
new file mode 100644
index 0000000..ce1636b
--- /dev/null
+++ b/bmcook/pruning/sprune/utils.py
@@ -0,0 +1,241 @@
+import types
+import torch
+from torch import Tensor
+from torch.nn import Module
+from typing import List, Dict, Union, Optional, NewType
+
+SPlugin = NewType('Plugin', Dict[str, Optional[Union[int, Tensor]]])
+
+################################# get config from model #################################
+def get_dim_ff(module: Module):
+    if hasattr(module, 'dim_ff') and isinstance(module.dim_ff, int):
+        ret = module.dim_ff
+    else:
+        ret = module.w_out.dim_in
+    return ret
+
+################################# utils for forward-inside pruning setup #################################
+def set_pruning_transformer(module: Module, index: int, TRANSFORMER_MASK: List[SPlugin], is_bmtCBlock: bool = True):
+    module.index = index
+    if is_bmtCBlock:
+        module.forward_unprune = module._module.forward
+    else:
+        module.forward_unprune = module.forward
+    def prune_forward(module_self, self_hidden_states, *args, **kwargs):
+        index = module_self.index
+        mask = TRANSFORMER_MASK[index]['mask']
+        out = module_self.forward_unprune(self_hidden_states, *args, **kwargs)
+        if mask is not None:
+            out = self_hidden_states + (out - self_hidden_states) * mask
+        return out
+    if is_bmtCBlock:
+        module._module.forward = types.MethodType(prune_forward, module)
+    else:
+        module.forward = types.MethodType(prune_forward, module)
+
+def set_pruning_att(module: Module, index: int, ATTENTION_MASK: List[SPlugin], NUM_HEADS_MASK: List[SPlugin], DIM_HEAD_MASK: List[SPlugin]):
+    module.index = index
+    module.forward_unprune = module.forward
+    def prune_forward(module_self, hidden_states, *args, **kwargs):
+        index = module_self.index
+        mask = ATTENTION_MASK[index]['mask']
+        out = module_self.forward_unprune(hidden_states, *args, **kwargs)
+        if mask is not None:
+            out = hidden_states + (out - hidden_states) * mask
+        return out
+    module.forward = types.MethodType(prune_forward, module)
+    # prune Linears:
+    for s_name, s_module in module.named_modules():
+        if 'project' in s_name:
+            set_pruning_linear_attention(s_module, index, NUM_HEADS_MASK, DIM_HEAD_MASK, 'in')
+        elif 'attention_out' in s_name:
+            set_pruning_linear_attention(s_module, index, NUM_HEADS_MASK, DIM_HEAD_MASK, 'out')
+
+def set_pruning_ffn(module: Module, index: int, FFN_MASK: List[SPlugin], DIM_FF_MASK: List[SPlugin]):
+    module.index = index
+    module.forward_unprune = module.forward
+    def prune_forward(module_self, hidden_states, *args, **kwargs):
+        index = module_self.index
+        mask = FFN_MASK[index]['mask']
+        out = module_self.forward_unprune(hidden_states, *args, **kwargs)
+        if mask is not None:
+            out = hidden_states + (out - hidden_states) * mask
+        return out
+    module.forward = types.MethodType(prune_forward, module)
+    # prune Linears
+    for s_name, s_module in module.named_modules():
+        if 'w_in.w' in s_name:
+            set_pruning_linear_feedforward(s_module, index, DIM_FF_MASK, 'in')
+        elif 'w_out' in s_name:
+            set_pruning_linear_feedforward(s_module, index, DIM_FF_MASK, 'out')
+
+def set_pruning_linear_attention(module: Module, index: int, NUM_HEADS_MASK: List[SPlugin], DIM_HEAD_MASK: List[SPlugin], in_out: str):
+    assert len(NUM_HEADS_MASK) == len(DIM_HEAD_MASK)
+    module.index = index
+    module.forward_unprune = module.forward
+    if in_out == 'in':
+        def prune_forward(module_self, *args, **kwargs):
+            index = module_self.index
+            num_heads, dim_head = NUM_HEADS_MASK[index]['dim'], DIM_HEAD_MASK[index]['dim']
+            num_heads_mask, dim_head_mask = NUM_HEADS_MASK[index]['mask'], DIM_HEAD_MASK[index]['mask']
+            out = module_self.forward_unprune(*args, **kwargs)  # (batch, len, num_heads*dim_head)
+            if num_heads_mask is not None:
+                old_size = out.size()
+                out = out.view(old_size[0], old_size[1], num_heads, dim_head)
+                out = out * num_heads_mask.view(num_heads, 1)
+                out = out.view(old_size[0], old_size[1], num_heads * dim_head)
+            if dim_head_mask is not None:
+                old_size = out.size()
+                out = out.view(old_size[0], old_size[1], num_heads, dim_head)
+                out = out * dim_head_mask
+                out = out.view(old_size[0], old_size[1], num_heads * dim_head)
+            return out
+    elif in_out == 'out':
+        def prune_forward(module_self, x, *args, **kwargs):
+            index = module_self.index
+            num_heads, dim_head = NUM_HEADS_MASK[index]['dim'], DIM_HEAD_MASK[index]['dim']
+            num_heads_mask, dim_head_mask = NUM_HEADS_MASK[index]['mask'], DIM_HEAD_MASK[index]['mask']
+            if num_heads_mask is not None:
+                old_size = x.size()  # (batch, len, num_heads * dim_head)
+                x = x.view(old_size[0], old_size[1], num_heads, dim_head)
+                x = x * num_heads_mask.view(num_heads, 1)
+                x = x.view(old_size[0], old_size[1], num_heads * dim_head)
+            if dim_head_mask is not None:
+                old_size = x.size()
+                x = x.view(old_size[0], old_size[1], num_heads, dim_head)
+                x = x * dim_head_mask
+                x = x.view(old_size[0], old_size[1], num_heads * dim_head)
+            out = module_self.forward_unprune(x, *args, **kwargs)  # (batch, len, dim_model)
+            return out
+    module.forward = types.MethodType(prune_forward, module)
+
+def set_pruning_linear_feedforward(module: Module, index: int, DIM_FF_MASK: List[SPlugin], in_out: str):
+    module.index = index
+    module.forward_unprune = module.forward
+    if in_out == 'in':
+        def prune_forward(module_self, *args, **kwargs):
+            index = module_self.index
+            mask = DIM_FF_MASK[index]['mask']
+            out = module_self.forward_unprune(*args, **kwargs)
+            if mask is not None:
+                out = out * mask
+            return out
+    elif in_out == 'out':
+        def prune_forward(module_self, x, *args, **kwargs):
+            index = module_self.index
+            mask = DIM_FF_MASK[index]['mask']
+            if mask is not None:
+                x = x * mask
+            out = module_self.forward_unprune(x, *args, **kwargs)
+            return out
+    module.forward = types.MethodType(prune_forward, module)
+
+
+################################# utils for cross-granularity sparsity computaion #################################
+def get_params_from_block(info_dict: Dict[int, Dict[str, List]]):
+    r"""The calculation of sparsity.
+
+    It calculates the mask and params of TransformerBlock, in a bottom-up way. This can support the cross-grained pruning.
+    """
+    res_exp, res_all = 0, 0
+    for k in info_dict:
+        modules = info_dict[k]['module']
+        params = info_dict[k]['param']
+        scores = info_dict[k]['score']
+
+        # transformer_score * (att_score * (num_heads_mask * dim_head_mask * 4) + cross_att_score * (cross_num_heads_mask * cross_dim_head_mask * 4) + ffn_score * (dim_ff_score * param * 3))
+        att_param_exp, cross_att_param_exp, ffn_param_exp = 0, 0, 0
+        att_param_all, cross_att_param_all, ffn_param_all = 0, 0, 0
+        if 'num_heads' in modules and 'dim_head' in modules:
+            num_heads_index = modules.index('num_heads')
+            num_heads_mask = scores[num_heads_index]
+            dim_head_index = modules.index('dim_head')
+            dim_head_mask = scores[dim_head_index]
+            att_param_exp = torch.sum(num_heads_mask, dtype=torch.float) * torch.sum(dim_head_mask) * 4
+            att_param_all = num_heads_mask.size(0) * dim_head_mask.size(0) * 4
+        elif 'num_heads' in modules:
+            num_heads_index = modules.index('num_heads')
+            num_heads_mask = scores[num_heads_index]
+            param = params[num_heads_index]
+            att_param_exp = torch.sum(num_heads_mask, dtype=torch.float) * param * 4
+            att_param_all = num_heads_mask.size(0) * param * 4
+        elif 'dim_head' in modules:
+            dim_head_index = modules.index('dim_head')
+            dim_head_mask = scores[dim_head_index]
+            param = params[dim_head_index]
+            att_param_exp = torch.sum(dim_head_mask, dtype=torch.float) * param * 4
+            att_param_all = dim_head_mask.size(0) * param * 4
+        
+        if 'cross_num_heads' in modules and 'cross_dim_head' in modules:
+            cross_num_heads_index = modules.index('cross_num_heads')
+            cross_num_heads_mask = scores[cross_num_heads_index]
+            cross_dim_head_index = modules.index('cross_dim_head')
+            cross_dim_head_mask = scores[cross_dim_head_index]
+            cross_att_param_exp = torch.sum(cross_num_heads_mask, dtype=torch.float) * torch.sum(cross_dim_head_mask, dtype=torch.float) * 4
+            cross_att_param_all = cross_num_heads_mask.size(0) * cross_dim_head_mask.size(0) * 4
+        elif 'cross_num_heads' in modules:
+            cross_num_heads_index = modules.index('cross_num_heads')
+            cross_num_heads_mask = scores[cross_num_heads_index]
+            param = params[cross_num_heads_index]
+            cross_att_param_exp = torch.sum(cross_num_heads_mask, dtype=torch.float) * param * 4
+            cross_att_param_all = cross_num_heads_mask.size(0) * param * 4
+        elif 'cross_dim_head' in modules:
+            cross_dim_head_index = modules.index('cross_dim_head')
+            cross_dim_head_mask = scores[cross_dim_head_index]
+            param = params[cross_dim_head_index]
+            cross_att_param_exp = torch.sum(cross_dim_head_mask, dtype=torch.float) * param * 4
+            cross_att_param_all = cross_dim_head_mask.size(0) * param * 4
+                
+        if 'dim_ff' in modules:
+            dim_ff_index = modules.index('dim_ff')
+            dim_ff_mask = scores[dim_ff_index]
+            dim_ff_param = params[dim_ff_index]
+            ffn_param_exp = torch.sum(dim_ff_mask, dtype=torch.float) * dim_ff_param * 3
+            ffn_param_all = dim_ff_mask.size(0) * dim_ff_param * 3
+
+        if 'att' in modules:
+            att_index = modules.index('att')
+            att_mask = scores[att_index]
+            param = params[att_index]
+            if att_param_exp == 0:
+                att_param_exp = param * att_mask
+            else:
+                att_param_exp = (att_param_exp + param - att_param_all) * att_mask
+            att_param_all = param
+        
+        if 'cross' in modules:
+            cross_att_index = modules.index('cross_att')
+            cross_att_mask = scores[cross_att_index]
+            param = params[cross_att_index]
+            if cross_att_param_exp == 0:
+                cross_att_param_exp = param * cross_att_mask
+            else:
+                cross_att_param_exp = (cross_att_param_exp + param - cross_att_param_all) * att_mask
+            cross_att_param_all = param
+        
+        if 'ffn' in modules:
+            ffn_index = modules.index('ffn')
+            ffn_mask = scores[ffn_index]
+            param = params[ffn_index]
+            if ffn_param_exp == 0:
+                ffn_param_exp = param * ffn_mask
+            else:
+                ffn_param_exp = (ffn_param_exp + param - ffn_param_all) * att_mask
+            ffn_param_all = param
+        
+        param_exp = att_param_exp + cross_att_param_exp + ffn_param_exp
+        param_all = att_param_all + cross_att_param_all + ffn_param_all
+        
+        if 'transformer' in modules:
+            transformer_index = modules.index('transformer')
+            transformer_mask = scores[transformer_index]
+            param = params[transformer_index]
+            if param_exp == 0:
+                param_exp = param * transformer_mask
+                param_all = param
+            else:
+                param_exp = param_exp * transformer_mask
+
+        res_exp += param_exp
+        res_all += param_all
+    return 1 - res_exp / res_all
\ No newline at end of file
diff --git a/bmcook/quant/__init__.py b/bmcook/quant/__init__.py
index e1fa3ef..ebb9f80 100644
--- a/bmcook/quant/__init__.py
+++ b/bmcook/quant/__init__.py
@@ -9,13 +9,16 @@ class BMQuant:
     '''
 
     @classmethod
-    def quantize(cls, model, config):
+    def quantize(cls, model, config, target_linear = model_center.layer.Linear):
         '''
         Practitioners can turn on quantization by `is_quant` in the config, which will replace all linear layers with quantized linear layers. BMCook provides the simulation of 8-bit quantization.
 
         :param model: Model to quantize.
         :param config: Configuration of the quantization.
         '''
+        quant_config = config.get('quantization')
+        if not quant_config['is_quant']:
+            return
 
         # fix cpm_kernel
         ct.gemm.GEMMInt8._backward = ct.gemm.GEMMInt8.backward
@@ -25,16 +28,16 @@ def new_func(ctx, grad_f):
             return ct.gemm.GEMMInt8._backward(ctx, grad_f)
         ct.gemm.GEMMInt8.backward = new_func
 
-        quant_config = config.get('quantization')
-        if not quant_config['is_quant']:
-            return
-
         for name, module in model.named_modules():
-            if isinstance(module, model_center.layer.Linear):
+            if isinstance(module, target_linear):
                 if len(quant_config["quantized_module"]) != 0:
                     if not any([pattern in name for pattern in quant_config["quantized_module"]]):
                         continue
-                module.forward = types.MethodType(forward_in8, module)
+                if target_linear != model_center.layer.Linear:
+                    module.forward = types.MethodType(forward_int8_cpmlive, module)
+                else:
+                    module.forward = types.MethodType(forward_in8, module)
+                module.quant = True
 
 def forward_in8(module_self, x):
     if module_self.length_scale and module_self.length_scale_before:
@@ -47,3 +50,13 @@ def forward_in8(module_self, x):
     if module_self.bias is not None:
         x = x + module_self.bias
     return x
+
+def forward_int8_cpmlive(module_self, x):
+    if module_self.scale_before:
+        x = x / math.sqrt(module_self.dim_in)
+    x = x.transpose(1, 2).contiguous()
+    x = ct.bmm(module_self.weight.unsqueeze(0), False, x, False, int8=True)
+    x = x.transpose(1, 2).contiguous()
+    if not module_self.scale_before:
+        x = x / math.sqrt(module_self.dim_in)
+    return x
\ No newline at end of file
diff --git a/bmcook/store.py b/bmcook/store.py
new file mode 100644
index 0000000..77ea04a
--- /dev/null
+++ b/bmcook/store.py
@@ -0,0 +1,168 @@
+import torch
+import bmtrain as bmt
+
+from collections import OrderedDict
+
+def _save_to_state_dict(model : torch.nn.Module, destination, prefix):
+    if isinstance(model, bmt.CheckpointBlock):
+        if bmt.global_var.config['rank'] != 0:
+            destination = OrderedDict() # creates an temporary ordered dict
+            destination._metadata = OrderedDict()
+        model.state_dict(destination, prefix, False)
+    else:
+        if bmt.global_var.config['rank'] != 0:
+            destination = OrderedDict() # creates an temporary ordered dict
+            destination._metadata = OrderedDict()
+        model._save_to_state_dict(destination, prefix, False)
+
+def _save_to_rank0(model : torch.nn.Module, destination=None, prefix=''):
+    if destination is None:
+        destination = OrderedDict()
+        destination._metadata = OrderedDict()
+    destination._metadata[prefix[:-1]] = local_metadata = dict(version=model._version)
+    _save_to_state_dict(model, destination, prefix)
+    for name, module in model._modules.items():
+        if module is not None:
+            _save_to_rank0(module, destination, prefix + name + '.')
+    for hook in model._state_dict_hooks.values():
+        hook_result = hook(model, destination, prefix, local_metadata)
+        if hook_result is not None:
+            destination = hook_result
+    return destination
+
+def save_spruned(model, plugin, file_name):
+    # get state_dict of model
+    torch.cuda.synchronize()
+    state_dict = _save_to_rank0(model)
+
+    new_state_dict = state_dict.copy()
+    config = model.config
+    
+    # reformat masks
+    num_encoder = plugin.info_to_engine['num_encoder_layers']
+    att_boundary = plugin.info_to_engine['att_boundary']
+    prune_layer = []
+    num_heads_mask, dim_head_mask, dim_ff_mask = {}, {}, {}
+    for mask in plugin.transformer:
+        if mask['mask'] == 0:
+            layer_index = mask['index']
+            if layer_index < num_encoder:
+                prune_layer.append('encoder.layers.' + str(layer_index))
+            else:
+                prune_layer.append('decoder.layers.' + str(layer_index - num_encoder))
+    
+    for i, mask in enumerate(plugin.att):
+        layer_index = mask['index']
+        if layer_index < num_encoder:
+            layer_prefix = 'encoder.layers.' + str(layer_index) + '.self_att'
+        else:
+            if i < att_boundary:
+                layer_prefix = 'decoder.layers.' + str(layer_index - num_encoder) + '.self_att'
+            else:
+                layer_prefix = 'decoder.layers.' + str(layer_index - num_encoder) + '.cross_att'
+        if mask['mask'] == 0:       # the att_layer pruned
+            prune_layer.append(layer_prefix)
+        else:
+            if plugin.num_heads[i]['mask'] is not None:
+                num_heads_mask[layer_prefix + '.self_attention'] = plugin.num_heads[i]['mask']
+            if plugin.dim_head[i]['mask'] is not None:
+                dim_head_mask[layer_prefix + '.self_attention'] = plugin.dim_head[i]['mask']
+    
+    for j, mask in enumerate(plugin.ffn):
+        layer_index = mask['index']
+        if layer_index < num_encoder:
+            layer_prefix = 'encoder.layers.' + str(layer_index) + '.ffn'
+        else:
+            layer_prefix = 'decoder.layers.' + str(layer_index - num_encoder) + '.ffn'
+        if mask['mask'] == 0:
+            prune_layer.append(layer_prefix)
+        elif plugin.dim_ff[j]['mask'] is not None:
+            dim_ff_mask[layer_prefix + '.ffn'] = plugin.dim_ff[j]['mask']
+
+    # do pruning
+    find_first_head = None
+    for k, v in state_dict.items():
+        if any([k.startswith(prefix) for prefix in prune_layer]):
+            del new_state_dict[k]
+    
+        elif 'encoder.layers' in k or 'decoder.layers' in k:
+            if 'project_' in k and k.split('.project_')[0] in num_heads_mask:              
+                k_prefix = k.split('.project_')[0]
+
+                if find_first_head is None:
+                    find_first_head  = num_heads_mask[k_prefix]
+                
+                heads_mask = num_heads_mask[k_prefix]['mask']
+                dim_head = num_heads_mask[k_prefix]['param']
+                num_heads_unprune = num_heads_mask[k_prefix]['dim']
+                num_heads_target = torch.sum(heads_mask).item()
+                v = v.view(num_heads_unprune, dim_head, config.dim_model)
+                tgt = []
+                for i, mask in enumerate(heads_mask):
+                    if mask == 1.:
+                        tgt.append(v[i])
+                v = torch.stack(tgt)
+                assert v.size() == (num_heads_target, dim_head, config.dim_model)
+                new_state_dict[k] = v.view(num_heads_target * dim_head, config.dim_model)
+            
+            elif 'attention_out' in k and k.split('.attention_out')[0] in num_heads_mask:
+                k_prefix = k.split('.attention_out')[0]
+
+                param_info = num_heads_mask[k_prefix]
+                heads_mask = param_info['mask']
+                dim_head = param_info['param']
+                num_heads_unprune = param_info['dim']
+                num_heads_target = torch.sum(heads_mask).item()
+                v = v.permute(1, 0).view(num_heads_unprune, dim_head, config.dim_model)
+                tgt = []
+                for i, mask in enumerate(heads_mask):
+                    if mask == 1.:
+                        tgt.append(v[i])
+                v = torch.stack(tgt)
+                assert v.size() == (num_heads_target, dim_head, config.dim_model)
+                new_state_dict[k] = v.view(num_heads_target * dim_head, config.dim_model).permute(1, 0)
+            
+            elif 'ffn.ffn.w_in' in k and k.split('.ffn.ffn.w_in')[0] in dim_ff_mask:
+                k_prefix = k.split('.ffn.ffn.w_in')[0]
+                
+                dimff_mask = dim_ff_mask[k_prefix]['mask']
+                dimff_unprune = dim_ff_mask[k_prefix]['dim']
+                dimff_target = torch.sum(dimff_mask).item()
+                assert v.size() == (dimff_unprune, config.dim_model)
+                tgt = []
+                for i, mask in enumerate(dimff_mask):
+                    if mask == 1.:
+                        tgt.append(v[i])
+                v = torch.stack(tgt)
+                assert v.size() == (dimff_target, config.dim_model)
+                new_state_dict[k] = v
+            
+            elif 'ffn.ffn.w_out' in k and k.split('.ffn.ffn.w_out')[0] in dim_ff_mask:
+                k_prefix = k.split('.ffn.ffn.w_out')[0]
+                
+                dimff_mask = dim_ff_mask[k_prefix]['mask']
+                dimff_unprune = dim_ff_mask[k_prefix]['dim']
+                dimff_target = torch.sum(dimff_mask).item()
+                v = v.permute(1, 0)
+                assert v.size() == (dimff_unprune, config.dim_model)
+                tgt = []
+                for i, mask in enumerate(dimff_mask):
+                    if mask == 1.:
+                        tgt.append(v[i])
+                v = torch.stack(tgt).permute(1, 0)
+                assert v.size() == (config.dim_model, dimff_target)
+                new_state_dict[k] = v
+        
+        if 'position_bias.relative_attention_bias' in k:
+            num_heads_unprune = find_first_head['dim']
+            heads_mask = find_first_head['mask']
+            v = v.permute(1, 0).view(num_heads_unprune, -1)
+            tgt = []
+            for i, mask in enumerate(heads_mask):
+                if mask == 1.:
+                    tgt.append(v[i])
+            v = torch.stack(tgt)
+            new_state_dict[k] = v.permute(1, 0)
+
+    if bmt.global_var.config["rank"] == 0:
+        torch.save(new_state_dict, file_name)
\ No newline at end of file
diff --git a/bmcook/trainer.py b/bmcook/trainer.py
new file mode 100644
index 0000000..f10cd98
--- /dev/null
+++ b/bmcook/trainer.py
@@ -0,0 +1,250 @@
+import types
+import inspect
+import subprocess
+import functools
+import bmtrain as bmt
+from typing import Optional, List
+from collections import OrderedDict
+from torch import Tensor
+from torch.nn import Module
+from torch.optim import Optimizer
+from packaging.version import parse as parse_version
+
+from .distilling import BMDistill
+from .pruning import BMPrune
+from .moe import BMMoE
+from .quant import BMQuant
+from .utils.config import ConfigParser
+
+class CookOutput:
+    def __init__(self, 
+        loss, 
+        original_output, 
+        lag_loss = 0., 
+        sparsity = 0., 
+        d_loss = 0., 
+        sprune_plugin = None, 
+        sprune_engine = None
+        ):
+
+        self.loss = loss
+        self.original_output = original_output
+        self.lag_loss = lag_loss
+        self.sparsity = sparsity
+        self.d_loss = d_loss
+        self.sprune_plugin = sprune_plugin
+        self.sprune_engine = sprune_engine
+
+def version_checker():
+    """
+    compatible with different ModelCenter version. (to help adjust the CookTrainer setup)
+
+    return True if the ModelCenter version is earlier than 0.1.3
+    """
+    result = subprocess.check_output('pip show model-center', shell=True)
+    version = parse_version(str(result).split('\\n')[1].split(' ')[-1])
+
+    return version <= parse_version('0.1.3')
+
+def remove_checkpointblock(model):
+    """remove CheckpointBlock in bmtrain, to get access to the forward func"""
+    for _, v in model.named_modules():
+
+        if isinstance(v, bmt.TransformerBlockList):
+
+            def new_func(list_self, hidden_states, *args):
+                for i in range(len(list_self._modules)):
+                    hidden_states = list_self._modules[str(i)](hidden_states, *args)
+                return hidden_states
+
+            v.forward = types.MethodType(new_func, v)
+
+            for k in v._modules.keys():
+                state_dict = v._modules[k].state_dict()
+                for kk, vv in v._modules[k]._module.named_modules():
+                    if kk+'.weight' in state_dict:
+                        vv.weight.data = state_dict[kk+'.weight'].clone().cuda()
+                    if kk+'.bias' in state_dict:
+                        vv.bias.data = state_dict[kk+'.bias'].clone().cuda()
+                v._modules[k] = v._modules[k]._module
+    return model
+
+
+class CookTrainer:
+    r"""A basic training manager of BMCook.
+
+    In BMCook, the distillation and pruning will introduce distill_loss, lagrangian_loss and sparsity. 
+    They are necessary for gradient backpropagation. CookTrainer will combine the original model outputs
+    and these necessary variables and return to users together.
+
+    Args:
+        _is_old_modelcenter(`bool`): the model center version checker. After the version 0.1.5 of model-center,
+        the output is reformed as a  :class:`ModelOutput`. So CookTrainer will first check the model-center
+        version, then modify the outputs.
+    
+    Example::
+    >>> # setup the forward
+    >>> CookTrainer.set_forward(your_cook_config, your_model, your_optimizer, your_teacher_model)
+    >>> ...
+    >>> # use the forward
+    >>> outputs = CookTrainer.forward(your_model, your_inputs, your_loss_func)
+    """
+
+    _is_old_modelcenter: bool = version_checker()
+
+    @staticmethod
+    def forward(model: Module, loss_func: Module, targets: Tensor, model_args: List, model_kwargs: OrderedDict) -> List:
+        r"""The core method in :class:`CookTrainer`.
+
+        it combines the original model.forward and loss calculation.
+
+        it should be rewrite in :method:`CookTrainer.set_forward`, if not an AttributeError will be raised.
+
+        Args:
+            model: target compressing model, basically from ModelCenter PLMs.
+            loss_func: the loss function. supports the format: loss_func(input: Tensor, target: Tensor) -> Tensor
+            targets: the target used to calculate loss.
+            model_args: args of :method:`model.forward`
+            model_kwargs: kwargs of :method:`model.forward`
+
+        Return:
+            `[loss, model_outputs, lag_loss, sparsity, d_loss, moe_hidden]`
+        """
+        raise NotImplementedError("The staticmethod forward() should be defined in :method:`set_forward`.")
+
+    @classmethod
+    def set_compression(cls, cook_config: ConfigParser, model: Optional[Module] = None, optimizer: Optional[Optimizer] = None, teacher: Optional[Module] = None, remove_ckptblock: bool = True):
+        r"""Define the :method:`forward`, and set up :class:`BMPrune`, :class:`BMDistill`, :class:`BMQuant`
+        and :class:`BMMoE`.
+
+        :class:`BMPrune` should be ahead of :class:`BMDistill`, because the :method:`forward` is 
+        chenged in both :class:`BMPrune` and :class:`BMDistill`.
+
+        The output format: `[loss, logits, lag_loss, sparsity, distill_loss]`
+
+        Args:
+            cook_config: BMCook config. You can set the compression strategy in this config and use 
+                `bmcook.utils.config.ConfigParse` to parse it.
+            model: target compressing model, basically from ModelCenter PLMs.
+            optimizer: optimizer used to train model.
+            teacher: teacher model used to distillation, basically from ModelCenter PLMs.
+        """
+        assert model is not None, "model is necessary for any compression setup."
+        
+        model_args_list = inspect.getfullargspec(model.forward).args
+        if teacher is not None:
+            teacher_args_list = inspect.getfullargspec(teacher.forward).args
+            if teacher_args_list != model_args_list:
+                raise ValueError("the techer forward func differs from model forward func.")
+        
+        if 'return_logits' in model_args_list:
+            model.forward = functools.partial(model.forward, return_logits=True)
+        elif 'output_logits' in model_args_list:
+            model.forward = functools.partial(model.forward, output_logits=True)
+
+        if cls._is_old_modelcenter:
+            def forward(model, loss_func, targets, *model_args, **model_kwargs):
+                outputs = model(
+                    *model_args, **model_kwargs)
+                logits = outputs
+                batch, seq_len, vocab_out_size = logits.size()
+
+                loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))
+                
+                ret = CookOutput(loss, outputs)
+                return ret
+        else:
+            def forward(model, loss_func, targets, *model_args, **model_kwargs):
+                outputs = model(
+                    *model_args, **model_kwargs)
+                logits = outputs.logits
+                batch, seq_len, vocab_out_size = logits.size()
+
+                loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))
+
+                ret = CookOutput(loss, outputs)
+                return ret
+
+        forward_doc = cls.forward.__doc__
+        cls.forward = forward
+        cls.forward.__doc__ = forward_doc
+
+        # remove CheckpointBlock
+        if remove_ckptblock:
+            model = remove_checkpointblock(model)
+
+        # for pruning
+        BMPrune.version = cls._is_old_modelcenter
+        BMPrune.compute_mask(model, cook_config)
+        cls.forward = BMPrune.set_forward_sprune(cls.forward)
+        BMPrune.set_optim_for_pruning(optimizer)
+
+        # for distillation
+        BMDistill.version = cls._is_old_modelcenter
+        cls.forward = BMDistill.set_forward(model, teacher, cls.forward, cook_config)
+
+        # for quantization
+        BMQuant.quantize(model, cook_config)
+
+        # for moefication
+        cls.forward = BMMoE.get_hidden(model, cook_config, cls.forward)
+
+        bmt.synchronize()
+
+
+class CPMAntTrainer:
+    r"""CookTrainer for CPM-Ant"""
+    
+    _is_old_modelcenter = version_checker()
+
+    @staticmethod
+    def forward(model, loss_func, targets, *model_args, **model_kwargs):
+        raise AttributeError("The staticmethod forward() should be defined in :method:`set_forward`.")
+    
+    @classmethod
+    def set_compression(cls, cook_config, model, optimizer, teacher, remove_ckptblock: bool = True, target_linear = None):
+        # remove CheckpointBlock
+
+        def forward(model, loss_func, targets, *model_args, **model_kwargs):
+            outputs = model(
+                *model_args, **model_kwargs)
+            logits = outputs[0]
+            # batch, seq_len, vocab_out_size = logits.size()
+
+            # loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))
+            loss = loss_func(logits.view(-1, logits.size(-1)), targets.view(-1))
+
+            ret = CookOutput(loss, outputs)
+            
+            return ret
+        forward_doc = cls.forward.__doc__
+        cls.forward = forward
+        cls.forward.__doc__ = forward_doc
+
+        # for quantization
+        if target_linear is not None:
+            BMQuant.quantize(model, cook_config, target_linear)
+        else:
+            BMQuant.quantize(model, cook_config)
+
+        # for pruning
+        BMPrune.version = cls._is_old_modelcenter
+        BMPrune.compute_mask(model, cook_config)
+        cls.forward = BMPrune.set_forward_sprune(cls.forward)
+        BMPrune.set_optim_for_pruning(optimizer)
+
+        # remove CheckpointBlock
+        if remove_ckptblock:
+            model = remove_checkpointblock(model)
+
+        # for distillation
+        BMDistill.version = cls._is_old_modelcenter
+        if target_linear is not None:
+            cls.forward = BMDistill.set_forward(model, teacher, cls.forward, cook_config, target_linear)
+        else:
+            cls.forward = BMDistill.set_forward(model, teacher, cls.forward, cook_config)
+
+        # for moefication
+        cls.forward = BMMoE.get_hidden(model, cook_config, cls.forward)
+
+        bmt.synchronize()
\ No newline at end of file
diff --git a/bmcook/utils/config.py b/bmcook/utils/config.py
index 64fa80a..36e5d2d 100644
--- a/bmcook/utils/config.py
+++ b/bmcook/utils/config.py
@@ -16,7 +16,14 @@
     "is_pruning": False,
     "pruning_mask_path": None,
     "pruned_module": ['[placehold]'],
-    "mask_method": "m4n2_1d/m4n2_2d",
+    "mask_method": "m4n2_1d/m4n2_2d/sprune",
+    "sprune": {
+        "criterion": "l0",
+        "training_mask": ['[placehold]'],
+        "fixed_mask_path": "",
+        "mask_mode": "train_mask",
+        "target_sparsity": 0.5
+    }
   },
 
   "quantization": {
@@ -27,7 +34,7 @@
   "MoEfication": {
     "is_moefy": False,
     "first_FFN_module": ['[placehold]'],
-  }  
+  }
 }
 
 class ConfigParser:
diff --git a/cpm_live_example/config/cpm-ant-1b.json b/cpm_live_example/config/cpm-ant-1b.json
deleted file mode 100644
index cbef433..0000000
--- a/cpm_live_example/config/cpm-ant-1b.json
+++ /dev/null
@@ -1,17 +0,0 @@
-{
-    "vocab_size": 30720, 
-    "dim_model": 4096, 
-    "dim_ff": 700, 
-    "num_layers": 48, 
-    "num_heads": 16, 
-    "dim_head": 128, 
-    "dropout_p": 0.0, 
-    "position_bias_num_buckets": 512, 
-    "position_bias_max_distance": 2048, 
-    "eps": 1e-06, 
-    "half": true,  
-    "prompt_types": 32, 
-    "prompt_length": 32, 
-    "segment_types": 32, 
-    "mask_modules": [[false, false], [false, true], [false, true], [true, true], [false, true], [false, false], [false, true], [false, false], [false, true], [false, false], [true, true], [false, false], [false, true], [true, true], [true, true], [true, true], [false, false], [true, false], [false, false], [true, true], [false, false], [true, true], [true, false], [false, true], [false, false], [true, true], [true, true], [false, false], [true, false], [false, true], [false, true], [false, false], [true, true], [false, true], [true, true], [true, true], [false, false], [true, false], [true, true], [false, false], [false, false], [true, true], [false, true], [true, true], [true, false], [true, false], [true, true], [true, false]]
-}
\ No newline at end of file
diff --git a/cpm_live_example/config/l0_pruning.json b/cpm_live_example/config/l0_pruning.json
deleted file mode 100644
index 50720a5..0000000
--- a/cpm_live_example/config/l0_pruning.json
+++ /dev/null
@@ -1,19 +0,0 @@
-{
-    "coarse-grained": {
-        "train_mask": false,
-        "coarse_mask": "results/sprune/2B_layer_mask.pt",
-        "target_sparsity": 0.67
-    },
-
-    "fine-grained": {
-        "train_mask": true,
-        "coarse_mask": "results/sprune/2B_layer_mask.pt",
-        "num_heads": 16,
-        "dim_ff": 700,
-        "target_num_heads": 4,
-        "target_dimff": 350,
-        "heads_mask": "results/sprune/300M_head_mask.pt",
-        "dimff_mask": "results/sprune/300M_dimff_mask.pt",
-        "prune_mode": "att"
-    }
-}
\ No newline at end of file
diff --git a/cpm_live_example/cpm_live/__init__.py b/cpm_live_example/cpm_live/__init__.py
deleted file mode 100644
index 03e328b..0000000
--- a/cpm_live_example/cpm_live/__init__.py
+++ /dev/null
@@ -1,5 +0,0 @@
-from . import models
-from . import dataset
-from . import utils
-from . import tokenizers
-from .arguments import get_args
diff --git a/cpm_live_example/cpm_live/arguments.py b/cpm_live_example/cpm_live/arguments.py
deleted file mode 100644
index 79f29c9..0000000
--- a/cpm_live_example/cpm_live/arguments.py
+++ /dev/null
@@ -1,150 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The OpenBMB team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import argparse
-
-
-def add_model_config_args(parser: argparse.ArgumentParser):
-    """Model arguments"""
-
-    group = parser.add_argument_group("model", "model configuration")
-    group.add_argument("--model-config", type=str, help="model configuration file")
-    return parser
-
-
-def add_training_args(parser: argparse.ArgumentParser):
-    """Training arguments."""
-
-    group = parser.add_argument_group("train", "training configurations")
-
-    group.add_argument(
-        "--base-path",
-        type=str,
-        default=None,
-        help="Path to the project base directory.",
-    )
-    group.add_argument("--dataset_name", type=str, default=None, help="Name of the dataset")
-    group.add_argument(
-        "--load",
-        type=str,
-        default=None,
-        help="Path to a directory containing a model checkpoint.",
-    )
-    group.add_argument(
-        "--save",
-        type=str,
-        default=None,
-        help="Output directory to save checkpoints to.",
-    )
-    group.add_argument(
-        "--save-name",
-        type=str,
-        default=None,
-        help="Output filename to save checkpoints to.",
-    )
-    group.add_argument(
-        "--save-iters",
-        type=int,
-        default=1000,
-        help="number of iterations between saves",
-    )
-    group.add_argument(
-        "--log-dir",
-        type=str,
-        default="logs",
-        help="tensorboard log directory",
-    )
-    group.add_argument("--inspect-iters", type=int, default=1000, help="number of inspecting")
-    group.add_argument("--batch-size", type=int, default=32, help="Data Loader batch size")
-    group.add_argument("--clip-grad", type=float, default=1.0, help="gradient clipping")
-    group.add_argument(
-        "--train-iters",
-        type=int,
-        default=1000000,
-        help="total number of iterations to train over all training runs",
-    )
-    group.add_argument("--max-length", type=int, default=512, help="max length of input")
-    group.add_argument(
-        "--max-encoder-length",
-        type=int,
-        default=512,
-        help="max length of encoder input",
-    )
-    group.add_argument(
-        "--max-decoder-length",
-        type=int,
-        default=256,
-        help="max length of decoder input",
-    )
-    group.add_argument(
-        "--start-step", type=int, default=0, help="step to start or continue training"
-    )
-    group.add_argument("--seed", type=int, default=1234, help="random seed for reproducibility")
-
-    group.add_argument(
-        "--epochs",
-        type=int,
-        default=1,
-        help="total number of epochs to train over all training runs",
-    )
-
-    # Learning rate.
-    group.add_argument("--lr", type=float, default=1.0e-4, help="initial learning rate")
-    group.add_argument("--weight-decay", type=float, default=1.0e-2, help="weight-decay")
-    group.add_argument("--loss-scale", type=float, default=65536, help="loss scale")
-
-    group.add_argument(
-        "--warmup-iters",
-        type=float,
-        default=0.01,
-        help="percentage of data to warmup on (.01 = 1% of all " "training iters). Default 0.01",
-    )
-    group.add_argument(
-        "--lr-decay-iters",
-        type=int,
-        default=None,
-        help="number of iterations to decay LR over,"
-        " If None defaults to `--train-iters`*`--epochs`",
-    )
-    group.add_argument(
-        "--lr-decay-style",
-        type=str,
-        default="noam",
-        choices=["constant", "linear", "cosine", "exponential", "noam"],
-        help="learning rate decay function",
-    )
-    group.add_argument(
-        "--local_rank",
-        type=int,
-        default=None,
-        help="local rank passed from distributed launcher",
-    )
-    group.add_argument(
-        "--data-bin-path",
-        type=str,
-        default=None,
-        help="path of the binary dataset"
-    )
-
-    return parser
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser = add_model_config_args(parser)
-    parser = add_training_args(parser)
-
-    args = parser.parse_args()
-    return args
diff --git a/cpm_live_example/cpm_live/dataset/__init__.py b/cpm_live_example/cpm_live/dataset/__init__.py
deleted file mode 100644
index 44ecec3..0000000
--- a/cpm_live_example/cpm_live/dataset/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-from .distributed_dataset import DistributedDataset, SimpleDataset, build_dataset
-from .utils import shuffle_dataset, compact_dataset, mask_dataset, merge_dataset
diff --git a/cpm_live_example/cpm_live/dataset/distributed_dataset.py b/cpm_live_example/cpm_live/dataset/distributed_dataset.py
deleted file mode 100644
index d90c612..0000000
--- a/cpm_live_example/cpm_live/dataset/distributed_dataset.py
+++ /dev/null
@@ -1,523 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The OpenBMB team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import io
-import os
-import pickle
-from typing import List
-import torch
-import bisect
-import bmtrain as bmt
-
-import random
-import string
-
-
-def _random_string():
-    return "".join(random.choices(string.ascii_uppercase + string.digits, k=8))
-
-
-class FileInfo:
-    def __init__(
-        self,
-        file_name: str,
-        block_begin: int,
-        block_end: int,
-        nbytes: int,
-        nlines: int,
-        mask: bool = False,
-    ) -> None:
-        self.file_name = file_name
-        self.block_begin = block_begin
-        self.block_end = block_end
-        self.nbytes = nbytes
-        self.nlines = nlines
-        self.mask = mask
-
-    @classmethod
-    def _load_from_state(cls, version, data):
-        if version == 1:
-            file_name, block_begin, block_end, nbytes, nlines, mask = data
-            return cls(file_name, block_begin, block_end, nbytes, nlines, mask)
-        else:
-            raise RuntimeError("Unsupported version %d" % version)
-
-    def __reduce__(self):
-        return (
-            FileInfo._load_from_state,
-            (
-                1,
-                (
-                    self.file_name,
-                    self.block_begin,
-                    self.block_end,
-                    self.nbytes,
-                    self.nlines,
-                    self.mask,
-                ),
-            ),
-        )
-
-
-_MASK_VALUE = 0x7FFFFFFF
-_DEFAULT_BLOCK_SIZE = 16 << 20
-
-
-class DistributedDataset:
-    """Open dataset in readonly mode.
-
-    `DistributeDataset` is used to read datasets in a distributed manner.
-    Data in this dataset will be distributed evenly in blocks to each worker in the `distributed communicator`.
-
-    **Note** When all data has been read, reading dataset again will revert back to the first data.
-
-    Args:
-        path (str): Path to dataset.
-        rank (int): Rank in distributed communicator. See: bmtrain.rank()
-        world_size (int): Total workers in distributed communicator. See: bmtrain.world_size()
-        block_size (int): Size of each block in bytes. All files in the same dataset should have the same block size. Default: 16MB
-
-    Example:
-        >>> dataset = DistributedDataset("/path/to/dataset")
-        >>> for i in range(10):
-        >>>     dataset.read()
-    """  # noqa: E501
-
-    def __init__(
-        self,
-        path: str,
-        rank: int = 0,
-        world_size: int = 1,
-        block_size=_DEFAULT_BLOCK_SIZE,
-    ) -> None:
-        # config
-        self._path = path
-        self._rank = rank
-        self._world_size = world_size
-        self._block_size = block_size
-
-        # dataset meta
-        self._block_states = torch.tensor([], dtype=torch.int)
-        self._file_info: List[FileInfo] = []
-        self._file_ends: List[int] = []
-        self._total_blocks = 0
-        self._nbytes = 0
-        self._nlines = 0
-
-        # states
-        self._curr_block = None
-        self._fp = None
-
-        # cache
-        self._last_mod_time = 0
-        self._curr_fname = None
-
-        self._update_states()
-
-    def _update_states(self, fast_skip: bool = True):
-        meta_path = os.path.join(self._path, "meta.bin")
-
-        mod_time = os.stat(meta_path).st_mtime
-        if self._last_mod_time < mod_time:
-            # file changed
-            pass
-        else:
-            if fast_skip:
-                return
-
-        info: List[FileInfo] = []
-        if os.path.exists(meta_path):
-            with open(meta_path, "rb") as f:
-                info = pickle.load(f)
-
-        old_len = len(self._file_info)
-        if old_len > len(info):
-            raise RuntimeError("Dataset meta file: changed unexpectly")
-        for i in range(old_len):
-            if self._file_info[i].file_name != info[i].file_name:
-                raise RuntimeError("Dataset meta file: changed unexpectly")
-            if self._file_info[i].block_begin != info[i].block_begin:
-                raise RuntimeError("Dataset meta file: changed unexpectly")
-            if self._file_info[i].block_end != info[i].block_end:
-                raise RuntimeError("Dataset meta file: changed unexpectly")
-
-        if info[0].block_begin != 0:
-            raise RuntimeError("Dataset meta file: block error (0)")
-        for i in range(len(info) - 1):
-            if info[i].block_end != info[i + 1].block_begin:
-                raise RuntimeError("Dataset meta file: block error (%d)" % (i + 1))
-
-        if old_len == len(info) and fast_skip:
-            # fast skip
-            return
-
-        if len(info) > 0:
-            total_blocks = info[-1].block_end
-            self._nbytes = 0
-            self._nlines = 0
-            for v in info:
-                self._nbytes += v.nbytes
-                self._nlines += v.nlines
-        else:
-            total_blocks = 0
-            self._nbytes = 0
-            self._nlines = 0
-
-        if total_blocks > 0:
-            masks = torch.full(
-                (total_blocks,),
-                _MASK_VALUE,
-                dtype=torch.int,
-                device="cpu",
-                requires_grad=False,
-            )
-            masks[self._rank :: self._world_size] = 0
-            for v in info:
-                if v.mask or (not os.path.exists(self._get_file_path(v.file_name))):
-                    masks[v.block_begin : v.block_end] = _MASK_VALUE
-            new_block_states = torch.zeros(
-                total_blocks, dtype=torch.int, device="cpu", requires_grad=False
-            )
-            new_block_states[: self._block_states.size(0)] = self._block_states
-            new_block_states = torch.maximum(new_block_states, masks)
-
-            self._block_states = new_block_states
-
-            self._file_ends = []
-            for v in info:
-                self._file_ends.append(v.block_end)
-        else:
-            self._block_states = torch.tensor(
-                [], dtype=torch.int, device="cpu", requires_grad=False
-            )
-            self._file_ends = []
-        self._total_blocks = total_blocks
-        self._file_info = info
-
-        assert len(self._file_ends) == len(self._file_info)
-
-    def _mask_file(self, f: FileInfo):
-        masks = torch.full(
-            (self._total_blocks,), 0, dtype=torch.int, device="cpu", requires_grad=False
-        )
-        masks[f.block_begin : f.block_end] = _MASK_VALUE
-        self._block_states = torch.maximum(self._block_states, masks)
-
-    def _get_block_file(self, block_id: int):
-        # find block in which file
-        file_idx = bisect.bisect_right(self._file_ends, block_id)
-        return self._file_info[file_idx]
-
-    def _get_next_block(self):
-        self._update_states()
-        if self._block_states.size(0) == 0:
-            raise RuntimeError("Empty dataset")
-        mn_block: int = self._block_states.argmin().item()  # type: ignore
-        if self._block_states[mn_block].item() == _MASK_VALUE:
-            raise RuntimeError("Empty dataset")
-        self._block_states[mn_block] += 1
-        return mn_block
-
-    def state_dict(self):
-        """Returns a state dict representing the read states of the dataset.
-
-        Example:
-            >>> state = dataset.state_dict()
-            >>> dataset.load_state_dict(state)
-        """
-        self._update_states()
-        states = torch.where(
-            self._block_states == _MASK_VALUE,
-            torch.zeros(self._total_blocks, dtype=torch.int, device="cpu", requires_grad=False),
-            self._block_states,
-        )
-
-        if (self._fp is not None) and (self._curr_block is not None):
-            curr_block = self._curr_block
-            curr_f = self._get_block_file(curr_block)
-            inblock_offset = self._fp.tell() - (curr_block - curr_f.block_begin) * self._block_size
-        else:
-            curr_block = -1
-            inblock_offset = 0
-
-        with torch.no_grad():
-            if self._world_size > 1:
-                gpu_states = states.cuda()
-                gpu_block = torch.tensor([curr_block, inblock_offset], dtype=torch.long).cuda()
-                global_states = bmt.distributed.all_reduce(gpu_states, op="sum").cpu()
-                global_block = bmt.distributed.all_gather(gpu_block).cpu()
-                return {"states": global_states, "block": global_block}
-            else:
-                return {
-                    "states": states,
-                    "block": torch.tensor(
-                        [[curr_block, inblock_offset]], dtype=torch.long, device="cpu"
-                    ),
-                }
-
-    def load_state_dict(self, state, strict: bool = True):
-        """Load dataset state.
-
-        Args:
-            state (dict): dataset state dict.
-            strict (bool): If `strict` is True, world size needs to be the same as when exported.
-
-        Example:
-            >>> state = dataset.state_dict()
-            >>>
-        """
-
-        self._block_states = state["states"]
-        self._update_states(False)
-
-        if state["block"].size(0) != self._world_size:
-            if strict:
-                raise ValueError(
-                    "world_size changed (%d -> %d)" % (state["block"].size(0), self._world_size)
-                )
-            else:
-                self._curr_block = None
-                self._fp = None
-                self._curr_fname = None
-        else:
-            curr_block = state["block"][self._rank][0].item()
-            inblock_offset = state["block"][self._rank][1].item()
-
-            if curr_block == -1:
-                self._curr_block = None
-            else:
-                self._curr_block = curr_block
-                f_info = self._get_block_file(self._curr_block)
-                self._open_file(
-                    f_info.file_name,
-                    (self._curr_block - f_info.block_begin) * self._block_size + inblock_offset,
-                )
-        # end
-
-    def _get_file_path(self, fname):
-        return os.path.join(self._path, fname)
-
-    def _open_file(self, fname, offset):
-        if self._curr_fname != fname:
-            if self._fp is not None:
-                self._fp.close()
-                self._curr_fname = None
-            self._fp = open(self._get_file_path(fname), "rb")
-            self._curr_fname = fname
-        else:
-            assert self._fp is not None, "Unexpected error"
-
-        self._fp.seek(offset, io.SEEK_SET)  # move to block
-
-    def read(self):
-        """Read a piece of data from dataset.
-
-        Workers in different ranks will read different data.
-        """
-        if self._curr_block is None:
-            next_block_id = self._get_next_block()
-            f_info = self._get_block_file(next_block_id)
-            try:
-                self._open_file(
-                    f_info.file_name,
-                    (next_block_id - f_info.block_begin) * self._block_size,
-                )
-                self._curr_block = next_block_id
-            except FileNotFoundError:
-                self._mask_file(f_info)
-                return self.read()  # read again
-
-        if self._fp is None:
-            raise RuntimeError("Dataset is not initialized")
-
-        MAGIC = self._fp.read(1)
-        if MAGIC == b"\x1F":
-            # correct
-            return pickle.load(self._fp)
-        elif MAGIC == b"\x00":
-            # end of block
-            self._curr_block = None
-            return self.read()  # read next block
-        else:
-            raise ValueError("Invalid magic header")
-
-    @property
-    def nbytes(self):
-        return self._nbytes
-
-
-class SimpleDataset(DistributedDataset):
-    def __init__(self, path: str, block_size=_DEFAULT_BLOCK_SIZE) -> None:
-        super().__init__(path, 0, 1, block_size)
-
-    def _get_next_block(self):
-        self._update_states()
-        if self._block_states.size(0) == 0:
-            raise RuntimeError("Empty dataset")
-        mn_block: int = self._block_states.argmin().item()  # type: ignore
-        if self._block_states[mn_block].item() >= 1:
-            raise EOFError("no more data")
-        self._block_states[mn_block] += 1
-        return mn_block
-
-    def __iter__(self):
-        while True:
-            try:
-                data = self.read()
-            except EOFError:
-                break
-            yield data
-
-    def __len__(self):
-        return self._nlines
-
-
-class DatasetWriter:
-    def __init__(self, fname, block_size):
-        self._fname = fname
-        self._block_size = block_size
-        self._fp = open(self._fname, "wb")
-        self._inblock_offset = 0
-
-        self._nbytes = 0
-        self._nlines = 0
-        self._nblocks = 1
-
-    def write(self, data):
-        """Write a piece of data into dataset.
-
-        Args:
-            data (Any): Serialization will be done using pickle.
-
-        Example:
-            >>> writer.write( "anything you want" )
-
-        """
-        byte_data = pickle.dumps(data)
-        if self._inblock_offset + 2 + len(byte_data) > self._block_size:
-            self._fp.write(
-                b"\x00" * (self._block_size - self._inblock_offset)
-            )  # fill the remaining space with 0
-            self._inblock_offset = 0
-            self._nblocks += 1
-            # we go to the next block
-        if self._inblock_offset + 2 + len(byte_data) > self._block_size:
-            raise ValueError("data is larger than block size")
-
-        self._nbytes += len(byte_data)
-        self._nlines += 1
-
-        self._inblock_offset += 1 + len(byte_data)
-        self._fp.write(b"\x1F")
-        self._fp.write(byte_data)
-
-    @property
-    def nbytes(self):
-        return self._nbytes
-
-    @property
-    def nblocks(self):
-        return self._nblocks
-
-    @property
-    def nlines(self):
-        return self._nlines
-
-    def close(self):
-        if not self._fp.closed:
-            self._fp.write(b"\x00" * (self._block_size - self._inblock_offset))
-            self._fp.close()
-
-
-class DatasetBuilder:
-    def __init__(self, path: str, dbname: str, block_size=_DEFAULT_BLOCK_SIZE) -> None:
-        self._block_size = block_size
-        self._path = path
-        self._dbname = dbname
-
-        if not os.path.exists(self._path):
-            os.makedirs(self._path)
-
-        meta_path = os.path.join(self._path, "meta.bin")
-
-        info: List[FileInfo] = []
-        if os.path.exists(meta_path):
-            with open(meta_path, "rb") as f:
-                info = pickle.load(f)
-
-        for v in info:
-            if v.file_name == dbname:
-                raise ValueError("Dataset name exists")
-
-        self._db_path = os.path.join(self._path, self._dbname)
-        if os.path.exists(self._db_path):
-            raise ValueError("File exists `%s`" % self._db_path)
-
-    def __enter__(self):
-        self._writer = DatasetWriter(self._db_path, self._block_size)
-        return self._writer
-
-    def __exit__(self, exc_type, exc_value, exc_traceback):
-        if self._writer is None:
-            raise RuntimeError("Unexpected call to __exit__")
-
-        self._writer.close()
-        if exc_type is not None:
-            print("Error while writing file")
-            if os.path.exists(self._db_path):
-                os.unlink(self._db_path)
-        else:
-            meta_path = os.path.join(self._path, "meta.bin")
-            info: List[FileInfo] = []
-            if os.path.exists(meta_path):
-                with open(meta_path, "rb") as f:
-                    info = pickle.load(f)
-            last_block = 0
-            if len(info) > 0:
-                last_block = info[-1].block_end
-            info.append(
-                FileInfo(
-                    self._dbname,
-                    last_block,
-                    last_block + self._writer.nblocks,
-                    self._writer.nbytes,
-                    self._writer.nlines,
-                    False,
-                )
-            )
-
-            # atomic write to meta file
-            random_fname = os.path.join(self._path, ".meta.bin.%s" % _random_string())
-            with open(random_fname, "wb") as f:
-                pickle.dump(info, f)
-            os.rename(random_fname, meta_path)
-
-        self._writer = None
-
-
-def build_dataset(path: str, dbname: str, block_size: int = _DEFAULT_BLOCK_SIZE):
-    """Open the dataset in write mode and returns a writer.
-
-    Args:
-        path (str): Path to dataset.
-        dbname (str): The name of the file to which the data will be written. The `dbname` needs to be unique in this `dataset`.
-        block_size (int): Size of each block in bytes. All files in the same dataset should have the same block size. Default: 16MB
-
-    Example:
-        >>> with build_dataset("/path/to/dataset", "data_part_1") as writer:
-        >>>     for i in range(10):
-        >>>         writer.write( { "anything you want" } )
-    """  # noqa: E501
-    return DatasetBuilder(path, dbname, block_size)
diff --git a/cpm_live_example/cpm_live/dataset/utils.py b/cpm_live_example/cpm_live/dataset/utils.py
deleted file mode 100644
index c166a56..0000000
--- a/cpm_live_example/cpm_live/dataset/utils.py
+++ /dev/null
@@ -1,254 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The OpenBMB team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import os
-from typing import List
-from .distributed_dataset import (
-    SimpleDataset,
-    build_dataset,
-    _random_string,
-    _DEFAULT_BLOCK_SIZE,
-    FileInfo,
-)
-import random
-import pickle
-import shutil
-
-try:
-    from tqdm import tqdm
-
-    support_tqdm = True
-except ModuleNotFoundError:
-    support_tqdm = False
-
-
-_DEFAULT_SHUFFLE_BUCKET_SIZE = 1 << 30
-
-
-def shuffle_dataset(
-    path_src: str,
-    path_tgt: str,
-    block_size: int = _DEFAULT_BLOCK_SIZE,
-    bucket_size: int = _DEFAULT_SHUFFLE_BUCKET_SIZE,
-    progress_bar=False,
-):
-    """Shuffle one distributed datataset, write results to another dataset.
-
-    Args:
-        path_str (str): path to source dataset
-        path_tgt (str): path to write results
-        block_size (int): dataset block size (default: 16MB)
-        bucket_size (int): shuffle algorithm bucket size (default: 1GB)
-        progress_bar (bool): show progress bar
-
-    Example:
-        >>> shuffle_dataset("/path/to/source", "/path/to/output")
-    """
-
-    if progress_bar and not support_tqdm:
-        raise RuntimeError("Requires `tqdm` to enable progress bar.")
-
-    ds = SimpleDataset(path_src, block_size=block_size)
-    num_buckets = (ds.nbytes + bucket_size - 1) // bucket_size
-
-    tmp_files = [os.path.join(path_src, ".tmp.%s" % _random_string()) for _ in range(num_buckets)]
-
-    try:
-        # Step 1: write to bucket randomly
-        f_tmp = [open(fname, "wb") for fname in tmp_files]
-        try:
-            iterator = ds
-            if progress_bar:
-                iterator = tqdm(ds, desc="Shuffle step 1/2")
-            for data in iterator:
-                bucket_id = int(random.random() * num_buckets)
-                pickle.dump(data, f_tmp[bucket_id])  # write into a random bucket
-        finally:
-            # close all files
-            for fp in f_tmp:
-                if not fp.closed:
-                    fp.close()
-        f_tmp = []
-
-        # Step 2: shuffle inside bucket
-        with build_dataset(path_tgt, "%s.shuffle" % _random_string()) as writer:
-            iterator = tmp_files
-            if progress_bar:
-                iterator = tqdm(tmp_files, desc="Shuffle step 2/2")
-
-            for fname in iterator:
-                fp = open(fname, "rb")
-                data_in_bucket = []
-                while True:
-                    try:
-                        data_in_bucket.append(pickle.load(fp))
-                    except EOFError:
-                        break
-                random.shuffle(data_in_bucket)
-                for data in data_in_bucket:
-                    writer.write(data)
-                fp.close()
-                os.unlink(fname)
-    finally:
-        # cleanup
-        for fname in tmp_files:
-            if os.path.exists(fname):
-                os.unlink(fname)
-
-
-def compact_dataset(path: str):
-    """Compact the dataset, removes blocks which the files were deleted.
-
-    **Note** This may affect the existing dataset state dict.
-
-    Args:
-        path (str): path to dataset
-
-    Example:
-        >>> compact_dataset("/path/to/dataset")
-
-    """
-
-    meta_path = os.path.join(path, "meta.bin")
-
-    info: List[FileInfo] = []
-    if os.path.exists(meta_path):
-        with open(meta_path, "rb") as f:
-            info = pickle.load(f)
-    else:
-        raise ValueError("Dataset not exists")
-
-    nw_info = []
-    curr_block = 0
-    for v in info:
-        if not os.path.exists(v.file_name):
-            # file is deleted
-            pass
-        else:
-            num_file_block = v.block_end - v.block_begin
-            nw_info.append(
-                FileInfo(
-                    v.file_name,
-                    curr_block,
-                    curr_block + num_file_block,
-                    v.nbytes,
-                    v.nlines,
-                    v.mask,
-                )
-            )
-            curr_block += num_file_block
-
-    random_fname = os.path.join(path, ".meta.bin.%s" % _random_string())
-    with open(random_fname, "wb") as f:
-        pickle.dump(nw_info, f)
-    os.rename(random_fname, meta_path)
-
-
-def mask_dataset(path: str, dbname: str, mask: bool = True):
-    """Mask one file in dataset. Blocks in masked datasets won't be read later.
-
-    Args:
-        path (str): path to dataset
-        dbname (str): file name in this dataset which you want to mask
-        mask (bool): True for mask, False for unmask
-
-    Example:
-        >>> mask_dataset("/path/to/dataset", "data_part_1", mask=True)
-
-    """
-
-    meta_path = os.path.join(path, "meta.bin")
-
-    info: List[FileInfo] = []
-    if os.path.exists(meta_path):
-        with open(meta_path, "rb") as f:
-            info = pickle.load(f)
-    else:
-        raise ValueError("Dataset not exists")
-
-    for v in info:
-        if v.file_name == dbname:
-            v.mask = mask
-
-    random_fname = os.path.join(path, ".meta.bin.%s" % _random_string())
-    with open(random_fname, "wb") as f:
-        pickle.dump(info, f)
-    os.rename(random_fname, meta_path)
-
-
-def merge_dataset(dst: str, src: str):
-
-    meta_path_src = os.path.join(src, "meta.bin")
-    meta_path_dst = os.path.join(dst, "meta.bin")
-
-    info_src: List[FileInfo] = []
-    if os.path.exists(meta_path_src):
-        with open(meta_path_src, "rb") as f:
-            info_src = pickle.load(f)
-    else:
-        raise ValueError("Dataset not exists")
-
-    info_dst: List[FileInfo] = []
-    if os.path.exists(meta_path_dst):
-        with open(meta_path_dst, "rb") as f:
-            info_dst = pickle.load(f)
-    else:
-        raise ValueError("Dataset not exists")
-
-    curr_block = 0
-    nw_info: List[FileInfo] = []
-    for v in info_dst:
-        num_file_block = v.block_end - v.block_begin
-        nw_info.append(
-            FileInfo(
-                v.file_name,
-                curr_block,
-                curr_block + num_file_block,
-                v.nbytes,
-                v.nlines,
-                v.mask,
-            )
-        )
-        curr_block += num_file_block
-
-    for v in info_src:
-        num_file_block = v.block_end - v.block_begin
-
-        dst_db_name = os.path.join(dst, v.file_name)
-        nw_fname = v.file_name
-        if os.path.exists(dst_db_name):
-            idx = 0
-            while os.path.exists(dst_db_name + "_{}".format(idx)):
-                idx += 1
-            dst_db_name = dst_db_name + "_{}".format(idx)
-            nw_fname = nw_fname + "_{}".format(idx)
-
-        shutil.copy(os.path.join(src, v.file_name), dst_db_name)
-        nw_info.append(
-            FileInfo(
-                nw_fname,
-                curr_block,
-                curr_block + num_file_block,
-                v.nbytes,
-                v.nlines,
-                v.mask,
-            )
-        )
-        curr_block += num_file_block
-
-    random_fname = os.path.join(dst, ".meta.bin.%s" % _random_string())
-    with open(random_fname, "wb") as f:
-        pickle.dump(nw_info, f)
-    os.rename(random_fname, meta_path_dst)
diff --git a/cpm_live_example/cpm_live/layers/__init__.py b/cpm_live_example/cpm_live/layers/__init__.py
deleted file mode 100644
index d0e5781..0000000
--- a/cpm_live_example/cpm_live/layers/__init__.py
+++ /dev/null
@@ -1,8 +0,0 @@
-from .embedding import Embedding
-from .position_embedding import SegmentPositionEmbedding
-from .linear import Linear
-from .layernorm import LayerNorm
-from .prune_attention import Attention
-from .prune_feedforward import FeedForward
-from .prune_blocks import TransformerBlock
-from .prune_transformer import Encoder
diff --git a/cpm_live_example/prune_cpm_ant.py b/cpm_live_example/prune_cpm_ant.py
deleted file mode 100644
index 95a0f41..0000000
--- a/cpm_live_example/prune_cpm_ant.py
+++ /dev/null
@@ -1,475 +0,0 @@
-import time
-from typing import Optional
-import torch
-import bmtrain as bmt
-import json
-import os
-import datetime
-from cpm_live.dataset import DistributedDataset
-from cpm_live import get_args
-import distutils.version  # noqa: F401
-from torch.utils.tensorboard import SummaryWriter
-
-from cpm_live.models import CPMAnt, CPMAntConfig
-from cpm_live.tokenizers import CPMAntTokenizer
-from training_tasks.ant import CPMAntPretrainDataset
-
-from bmcook.pruning import BMPrune
-
-def get_tokenizer(args):
-    tokenizer = CPMAntTokenizer()
-    return tokenizer
-
-
-def get_model(args):
-    config = CPMAntConfig.from_json_file(args.model_config)
-    model = CPMAnt(config)
-    if args.load is not None:
-        bmt.load(model, args.load)
-    else:
-        bmt.init_parameters(model)
-    args.prompt_length = config.prompt_length
-    return model
-
-
-def get_optimizer(args, model):
-    optimizer = bmt.optim.AdamOffloadOptimizer(
-        model.parameters(), weight_decay=args.weight_decay, scale=args.loss_scale
-    )
-    if args.load is not None:
-        if os.path.exists(os.path.join(args.save, args.save_name + (".rank-%d.opt" % 0))):
-            # optimizer state exists
-            states = torch.load(
-                os.path.join(args.save, args.save_name + (".rank-%d.opt" % bmt.rank()))
-            )
-            optimizer.load_state_dict(states)
-    return optimizer
-
-
-def get_learning_rate_scheduler(args, optimizer):
-    if args.lr_decay_iters is None:
-        args.lr_decay_iters = args.train_iters * args.epochs
-    lr_scheduler = bmt.lr_scheduler.Noam(
-        optimizer,
-        start_lr=args.lr,
-        warmup_iter=args.warmup_iters,
-        end_iter=args.lr_decay_iters,
-        num_iter=args.start_step,
-    )
-    return lr_scheduler
-
-
-def setup_model_and_optimizer(args):
-    tokenizer = get_tokenizer(args)
-    model = get_model(args)
-    bmt.synchronize()
-    optimizer = get_optimizer(args, model)
-    lr_scheduler = get_learning_rate_scheduler(args, optimizer)
-    bmt.synchronize()
-    return tokenizer, model, optimizer, lr_scheduler
-
-
-def initialize():
-    args = get_args()
-    bmt.init_distributed(seed=args.seed, loss_scale_factor=2, loss_scale_steps=512)
-    if args.save is not None:
-        os.makedirs(args.save, exist_ok=True)
-    return args
-
-
-def see_memory(detail=False):
-    if detail:
-        res = torch.cuda.memory_summary()
-    else:
-        res = (
-            round(torch.cuda.memory_allocated() / (1024 * 1024 * 1024), 2),
-            round(torch.cuda.max_memory_allocated() / (1024 * 1024 * 1024), 2),
-        )
-    torch.cuda.reset_peak_memory_stats()
-    return res
-
-
-def add_mem_time(info, mem_usage, tim_usage):
-    torch.cuda.synchronize()
-    mem_usage[info] = see_memory()
-    tim_usage[info] = time.time()
-    return mem_usage, tim_usage
-
-
-def get_log_time() -> datetime.datetime:
-    return datetime.datetime.utcnow() + datetime.timedelta(hours=16)
-
-
-def get_log_name(now: Optional[datetime.datetime] = None) -> str:
-    if now is None:
-        now = get_log_time()
-    return "log.%s.txt" % now.strftime("%Y%m%d")
-
-
-def lookup_latest_log():
-    now = get_log_time()
-
-    # try to find the latest log in the last 15 days
-    for _ in range(15):
-        log_name = get_log_name(now)
-        if os.path.exists(log_name):
-            with open(log_name, "r") as flog:
-                line = flog.readlines()[-1]  # get last log
-                return json.loads(line)
-
-        now -= datetime.timedelta(days=1)  # try the previous day
-    return None
-
-
-def get_tasks():
-    TASK_FILE_NAME = "tasks.json"
-    if not hasattr(get_tasks, "info"):
-        m_time = os.stat(TASK_FILE_NAME).st_mtime
-        tasks = json.load(open(TASK_FILE_NAME, "r", encoding="utf-8"))
-        get_tasks.info = {"m_time": m_time, "tasks": tasks}
-    else:
-        m_time = os.stat(TASK_FILE_NAME).st_mtime
-        if m_time != get_tasks.info["m_time"]:
-            tasks = json.load(open(TASK_FILE_NAME, "r", encoding="utf-8"))
-            get_tasks.info = {"m_time": m_time, "tasks": tasks}
-    return get_tasks.info["tasks"]
-
-
-class BatchPacker:
-    def __init__(self, dataset, max_length, batch_size):
-        self.dataset = dataset
-        self.max_length = max_length
-        self.batch_size = batch_size
-
-    def __iter__(self):
-        ctx = []
-        tgt = []
-        context = []
-        position = []
-        segment = []
-        span = []
-        task_info = []
-
-        for data in self.dataset:
-            (
-                ctx_data,
-                tgt_data,
-                _len,
-                context_data,
-                position_data,
-                segment_data,
-                task_data,
-            ) = data
-            if ctx_data is None:
-                continue
-            assert _len <= self.max_length
-
-            ctx_data = ctx_data.astype("int64")
-            tgt_data = tgt_data.astype("int64")
-
-            for index in range(len(ctx)):
-                if span[index][-1] + _len < self.max_length:
-                    ctx[index][span[index][-1] : span[index][-1] + _len] = torch.from_numpy(
-                        ctx_data
-                    )[:_len].long()
-                    tgt[index][span[index][-1] : span[index][-1] + _len] = torch.from_numpy(
-                        tgt_data
-                    )[:_len].long()
-                    context[index][span[index][-1] : span[index][-1] + _len] = torch.from_numpy(
-                        context_data
-                    )[:_len].bool()
-                    position[index][span[index][-1] : span[index][-1] + _len] = torch.from_numpy(
-                        position_data
-                    )[:_len].long()
-                    segment[index][span[index][-1] : span[index][-1] + _len] = torch.from_numpy(
-                        segment_data
-                    )[:_len].long()
-                    task_info[index][span[index][-1] : span[index][-1] + _len] = torch.from_numpy(
-                        task_data
-                    )[:_len].long()
-                    span[index].append(span[index][-1] + _len)
-                    break
-            else:
-
-                _ctx = torch.zeros((self.max_length,), dtype=torch.long)
-                _ctx[:_len] = torch.from_numpy(ctx_data)[:_len].long()
-                _tgt = torch.full((self.max_length,), -100, dtype=torch.long)
-                _tgt[:_len] = torch.from_numpy(tgt_data)[:_len].long()
-                _context = torch.full((self.max_length,), False, dtype=torch.bool)
-                _context[:_len] = torch.from_numpy(context_data)[:_len].bool()
-                _position = torch.full((self.max_length,), False, dtype=torch.long)
-                _position[:_len] = torch.from_numpy(position_data)[:_len].long()
-                _segment = torch.full((self.max_length,), False, dtype=torch.long)
-                _segment[:_len] = torch.from_numpy(segment_data)[:_len].long()
-                _task_info = torch.full((self.max_length,), -1, dtype=torch.long)
-                _task_info[:_len] = torch.from_numpy(task_data)[:_len].long()
-                ctx.append(_ctx)
-                tgt.append(_tgt)
-                context.append(_context)
-                position.append(_position)
-                segment.append(_segment)
-                task_info.append(_task_info)
-                span.append([_len])
-
-            if len(ctx) > self.batch_size:
-                _span = torch.zeros((self.batch_size, self.max_length + 1), dtype=torch.long)
-                for bindex in range(self.batch_size):
-                    for sindex in span[bindex]:
-                        _span[bindex][sindex] = 1
-
-                yield {
-                    "ctx": torch.stack(ctx[: self.batch_size]),
-                    "tgt": torch.stack(tgt[: self.batch_size]),
-                    "context": torch.stack(context[: self.batch_size]),
-                    "segment": torch.stack(segment[: self.batch_size]),
-                    "position": torch.stack(position[: self.batch_size]),
-                    "span": torch.cumsum(_span, dim=-1)[:, :-1],
-                    "len_ctx": torch.LongTensor([it[-1] for it in span[: self.batch_size]]),
-                    "task": torch.stack(task_info[: self.batch_size]),
-                }
-
-                ctx = ctx[self.batch_size :]
-                tgt = tgt[self.batch_size :]
-                context = context[self.batch_size :]
-                segment = segment[self.batch_size :]
-                position = position[self.batch_size :]
-                span = span[self.batch_size :]
-                task_info = task_info[self.batch_size :]
-
-
-def pretrain(args, tokenizer, model, optimizer, lr_scheduler, dataset):
-
-    # bmcook
-    from bmcook.utils.config import ConfigParser
-    config = ConfigParser('config/bmcook.json')
-    if config.get('pruning')['is_pruning'] == False:
-        train_mask, masks = False, None
-    else:
-        train_mask, masks = BMPrune.compute_mask(model, config)
-    if train_mask:
-        sprune_module = BMPrune._sprune_module
-        (sp_optimizer, lag_optimizer) = BMPrune._optim_and_scheduler 
-        expected_sparsity = 0
-
-    average_time = bmt.utils.AverageRecorder()
-    loss_func = bmt.loss.FusedCrossEntropy(ignore_index=-100)
-
-    start_step = args.start_step
-
-    if bmt.rank() == 0:
-        writer = SummaryWriter(log_dir=args.log_dir)
-
-    global_token_pass = 0.0
-    global_throughout = 0.0
-    global_world_size = bmt.world_size()
-
-    if bmt.rank() == 0:
-        latest_log = lookup_latest_log()
-        if latest_log is not None:
-            global_token_pass = latest_log["token pass"]
-
-    if os.path.exists(os.path.join(args.save, args.save_name + ("-%d.data.pt" % start_step))):
-        # load dataset states if exists
-        dataset_states = torch.load(
-            os.path.join(args.save, args.save_name + ("-%d.data.pt" % start_step))
-        )
-        dataset.dataset.load_state_dict(dataset_states)
-
-    dataloader = BatchPacker(dataset, args.max_length, args.batch_size)
-
-    for iteration, data in enumerate(dataloader):
-        
-        if train_mask:
-            sprune_module.zero_grad()
-            sp_optimizer.zero_grad()
-            lag_optimizer.zero_grad()
-        
-            bmt.print_rank('sprune_module forward pass and update the mask...')
-            masks = sprune_module.forward(training=True)
-
-        iteration = iteration + start_step + 1
-        assert len(data["ctx"]) == args.batch_size
-        input_idx = data["ctx"].int().cuda()
-        input_length = data["len_ctx"].int().cuda()
-        input_context = data["context"].bool().cuda()
-        input_position = data["position"].float().cuda()
-        input_segment = data["segment"].int().cuda()
-        input_span = data["span"].int().cuda()
-        targets = data["tgt"].long().cuda()
-        task_info = data["task"].long().cuda()
-
-        # ===========
-        optimizer.zero_grad()
-        torch.cuda.empty_cache()
-        mem_usage = {}
-        tim_usage = {}
-        mem_usage, tim_usage = add_mem_time("init", mem_usage, tim_usage)
-
-        # ===========
-        logits, _ = model(
-            input_idx,
-            input_length,
-            input_context,
-            input_position,
-            input_segment,
-            input_span,
-            masks=masks,
-        )
-        loss = loss_func(logits.view(-1, logits.size(-1)), targets.view(-1))
-        global_loss = bmt.sum_loss(loss).item()
-        mem_usage, tim_usage = add_mem_time("forward", mem_usage, tim_usage)
-        if train_mask:
-            lag_loss, expected_sparsity, _  = sprune_module.lagrangian_regularization()
-            lag_loss.backward()
-
-        # ===========
-        loss = optimizer.loss_scale(loss)
-        loss.backward()
-        mem_usage, tim_usage = add_mem_time("backward", mem_usage, tim_usage)
-
-        # ===========
-        grad_norm = bmt.optim.clip_grad_norm(
-            optimizer.param_groups, args.clip_grad, scale=optimizer.scale, norm_type=2
-        )
-        bmt.optim_step(optimizer, lr_scheduler)
-        mem_usage, tim_usage = add_mem_time("optim", mem_usage, tim_usage)
-        if train_mask:
-            sp_optimizer.step()
-            lag_optimizer.step()
-
-        # ==========
-        iteration_time = tim_usage["optim"] - tim_usage["init"]
-        average_time.record(iteration_time)
-
-        task_ids = get_tasks()
-        with torch.no_grad():
-            task_num = len(task_ids)
-            logits_tmp = logits.view(-1, logits.size(-1)).expand(task_num, -1, -1)
-            targets_tmp = targets.expand(task_num, -1, -1)
-            task_info = task_info.expand(task_num, -1, -1)
-
-            task = task_info.new([x for x in range(task_num)])[:, None, None]
-            targets_tmp = torch.where(task_info == task, targets_tmp, -100)
-
-            task_loss_list = []
-            for i in range(task_num):
-                task_loss = loss_func(logits_tmp[i, :], targets_tmp[i, :].view(-1))
-                global_task_loss = bmt.gather_result(task_loss.unsqueeze(0)).nanmean().item()
-                task_loss_list.append(global_task_loss)
-
-        local_total_rate = torch.Tensor([input_length.float().mean() / args.max_length]).cuda()
-        local_total_rate = bmt.sum_loss(local_total_rate).item()
-        global_token_pass += (
-            global_world_size
-            * local_total_rate
-            * (args.max_length - args.prompt_length)
-            * args.batch_size
-        )
-        avg_time = average_time.value
-        global_throughout += (
-            (args.max_length - args.prompt_length) * args.batch_size * local_total_rate / avg_time
-        )
-
-        train_info = {
-            "time": tim_usage["init"],
-            "iter": iteration,
-            "loss": global_loss,
-            "lr": lr_scheduler.current_lr,
-            "lr scale": int(optimizer.scale),
-            "time usage": tim_usage,
-            "mem usage": mem_usage,
-            "avg time (s)": avg_time,
-            "token/max": local_total_rate,
-            "token pass": global_token_pass,
-            "throughout (token/s)": (args.max_length - args.prompt_length)
-            * args.batch_size
-            * local_total_rate
-            / avg_time,
-            "global throughout (token/s)": global_throughout / iteration,
-            "grad_norm": grad_norm.item(),
-            "mask/max": ((targets >= 0).sum(-1).float().mean() / args.max_length).item(),
-            "num_gpus": global_world_size,
-        }
-        task_loss = {task_name: task_loss_list[idx] for (task_name, idx) in task_ids.items()}
-        train_info["task_loss"] = task_loss
-
-        bmt.print_rank(
-            (
-                "| Iter: {:6d} | loss: {:.4f} | lr: {:.4e}, scale: {:10.4f} | time: {:.4f} |"
-                + " token/max: {:.4f} | mask/max: {:.4f} | grad_norm: {:.4f}"
-            ).format(
-                iteration,
-                global_loss,
-                lr_scheduler.current_lr,
-                int(optimizer.scale),
-                avg_time,
-                input_length.float().mean() / args.max_length,
-                (targets >= 0).sum(-1).float().mean() / args.max_length,
-                grad_norm,
-            )
-        )
-        bmt.print_rank(
-            "| "
-            + " | ".join(
-                [
-                    "{} loss: {:.4f}".format(task_name, task_loss_list[idx])
-                    for task_name, idx in task_ids.items()
-                ]
-            )
-        )
-        if iteration % args.inspect_iters == 0:
-            model_inspect = bmt.inspect.inspect_model(model, "*")
-            bmt.print_rank(bmt.inspect.format_summary(model_inspect))
-            train_info["model_inspect"] = model_inspect
-
-        if bmt.rank() == 0:
-
-            with open(get_log_name(), "a") as ff:
-                ff.write(json.dumps(train_info) + "\n")
-
-        if bmt.rank() == 0:
-            writer.add_scalar("Loss/train", global_loss, iteration)
-            for i in task_ids.keys():
-                writer.add_scalar("Loss/train/{}".format(i), task_loss_list[task_ids[i]], iteration)
-
-        if args.save is not None and iteration % args.save_iters == 0:
-            bmt.save(model, os.path.join(args.save, args.save_name + ("-%d.pt" % iteration)))
-            torch.save(
-                optimizer.state_dict(),
-                os.path.join(args.save, args.save_name + (".rank-%d.opt" % bmt.rank())),
-            )
-            all_states = dataset.dataset.state_dict()
-            if bmt.rank() == 0:
-                # rank 0 writes the dataloader state
-                torch.save(
-                    all_states,
-                    os.path.join(args.save, args.save_name + ("-%d.data.pt" % iteration)),
-                )
-            del all_states
-
-        if train_mask: 
-            bmt.print_rank('expected_sparsity: ', expected_sparsity)
-            if bmt.rank() == 0:
-                writer.add_scalar("expected_sparsity", torch.mean(expected_sparsity), iteration)
-            if iteration % 20 == 0:
-                masks = sprune_module.forward(training=False)
-                torch.save(sprune_module, 'results/sprune/300M_l0_module.pt')
-                torch.save(masks, 'results/sprune/300M_head_mask.pt')
-
-    bmt.save(model, os.path.join(args.save, args.save_name + ".pt"))
-
-
-def main():
-    args = initialize()
-    tokenizer, model, optimizer, lr_scheduler = setup_model_and_optimizer(args)
-    dataset = CPMAntPretrainDataset(
-        DistributedDataset(args.data_bin_path, bmt.rank(), bmt.world_size()),
-        max_length=args.max_length - args.prompt_length,
-        prompt_length=args.prompt_length,
-        tokenizer=tokenizer,
-    )
-    pretrain(args, tokenizer, model, optimizer, lr_scheduler, dataset)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/cpm_live_example/scripts/prune_cpm_ant.sh b/cpm_live_example/scripts/prune_cpm_ant.sh
deleted file mode 100644
index e2d3e75..0000000
--- a/cpm_live_example/scripts/prune_cpm_ant.sh
+++ /dev/null
@@ -1,33 +0,0 @@
-#! /bin/bash
-export CUDA_VISIBLE_DEVICES=0,1,2,3
-GPUS_PER_NODE=4
-
-NNODES=1
-MASTER_ADDR=localhost
-MASTER_PORT=2333
-
-OPTS=""
-OPTS+=" --model-config cpm_live_example/config/cpm-ant-1b.json"
-OPTS+=" --batch-size 32"
-OPTS+=" --train-iters 200000"
-OPTS+=" --save-iters 500"
-OPTS+=" --save-name cpm_live_checkpoint"
-OPTS+=" --max-length 512"
-OPTS+=" --save results/"
-OPTS+=" --lr 0.1"
-OPTS+=" --inspect-iters 100"
-OPTS+=" --warmup-iters 2000"
-OPTS+=" --lr-decay-style noam"
-OPTS+=" --weight-decay 0.01"
-OPTS+=" --clip-grad 1.0"
-OPTS+=" --loss-scale 1048576"
-OPTS+=" --start-step 0"
-OPTS+=" --log-dir logs/tensorboard/cpm_live_48_4096/"
-OPTS+=" --load cpm_live_compress/ckpt/1B/cpm_live_checkpoint_1B_pruned.pt"
-OPTS+=" --data_bin_path /path/to/your/dataset"
-
-CMD="torchrun --nnodes=${NNODES} --nproc_per_node=${GPUS_PER_NODE} --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} prune_cpm_ant.py ${OPTS}"
-
-echo ${CMD}
-$CMD
-
diff --git a/cpm_live_example/tasks.json b/cpm_live_example/tasks.json
deleted file mode 100644
index 8cd069d..0000000
--- a/cpm_live_example/tasks.json
+++ /dev/null
@@ -1,4 +0,0 @@
-{
-    "mlm": 0,
-    "lm": 1
-}
\ No newline at end of file
diff --git a/cpm_live_example/training_tasks/__init__.py b/cpm_live_example/training_tasks/__init__.py
deleted file mode 100644
index 0a502f1..0000000
--- a/cpm_live_example/training_tasks/__init__.py
+++ /dev/null
@@ -1,1 +0,0 @@
-from . import ant
diff --git a/cpm_live_example/training_tasks/ant/__init__.py b/cpm_live_example/training_tasks/ant/__init__.py
deleted file mode 100644
index 5b99609..0000000
--- a/cpm_live_example/training_tasks/ant/__init__.py
+++ /dev/null
@@ -1,1 +0,0 @@
-from .pretrain import CPMAntPretrainDataset
diff --git a/cpm_live_example/training_tasks/ant/pretrain.py b/cpm_live_example/training_tasks/ant/pretrain.py
deleted file mode 100644
index 812cd52..0000000
--- a/cpm_live_example/training_tasks/ant/pretrain.py
+++ /dev/null
@@ -1,121 +0,0 @@
-# coding=utf-8
-# Copyright 2020 The OpenBMB team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from builtins import isinstance
-import torch.utils.data as data
-import random
-import torch
-import numpy as np
-
-
-class CPMAntPretrainDataset(data.Dataset):
-    def __init__(self, ctx, max_length=1024, prompt_length=32, tokenizer=None):
-        self.ctx = ctx
-        self.max_length = max_length + prompt_length
-        self.prompt_length = prompt_length
-        self.tokenizer = tokenizer
-
-    def __len__(self):
-        return len(self.ctx)
-
-    @property
-    def dataset(self):
-        return self.ctx
-
-    def __get_item_data(self, raw_data):
-        if isinstance(raw_data, torch.Tensor):
-            raw_data = raw_data.numpy()
-
-        global_task = raw_data[0]
-        n_segment = raw_data[1]
-        len_info = n_segment * 3 + 2
-        segment_len = raw_data[2:len_info:3]
-        segment_type = raw_data[3:len_info:3]
-        segment_task = raw_data[4:len_info:3]
-        ctx = raw_data[len_info:]
-
-        if ctx.shape[0] > self.max_length - self.prompt_length:
-            return None, None, None, None, None, None, None
-        len_ctx = min(ctx.shape[0], self.max_length - self.prompt_length)
-
-        context_inp = np.full(len_ctx, True)
-        position_inp = np.arange(len_ctx, dtype=np.int64)
-        segment_inp = np.full(len_ctx, 0, dtype=np.int64)
-        task_inp = np.full(len_ctx, 0, dtype=np.int64)
-        tgt = np.full(len_ctx, -100, dtype=np.int64)
-
-        # for each segment
-        segment_begin = 0
-        for i in range(n_segment):
-            segment_end = segment_begin + segment_len[i]
-            task = segment_task[i]
-            # generate target
-            if task == 0:
-                num_mask = random.randint(1, segment_len[i] - 1)
-                mask_idx = (
-                    np.random.choice(segment_len[i] - 1, num_mask, replace=False) + segment_begin
-                )
-                context_inp[mask_idx + 1] = False
-                assert segment_type[i] == 1
-            elif task == 1:
-                num_mask = random.randint(1, segment_len[i] - 1)
-                context_inp[segment_end - num_mask : segment_end] = False
-                assert segment_type[i] == 2
-            elif task == 3:
-                if segment_type[i] == 2:
-                    context_inp[1:] = False
-            elif task == 4:
-                if segment_type[i] == 3:
-                    context_inp[1:] = False
-            task_inp[segment_begin:segment_end] = task
-            segment_inp[segment_begin:segment_end] = segment_type[i]
-            tgt[segment_begin : segment_end - 1] = np.where(
-                context_inp[segment_begin + 1 : segment_end],
-                -100,
-                ctx[segment_begin + 1 : segment_end],
-            )
-            segment_begin = segment_end
-        # prepend prompt segment
-        context_inp = np.concatenate((np.full(self.prompt_length, True), context_inp))
-        position_inp = np.concatenate(
-            (
-                np.arange(self.prompt_length, dtype=np.int64),
-                position_inp + self.prompt_length,
-            )
-        )
-        segment_inp = np.concatenate((np.full(self.prompt_length, 0, dtype=np.int64), segment_inp))
-        task_inp = np.concatenate((np.full(self.prompt_length, 0, dtype=np.int64), task_inp))
-        tgt = np.concatenate((np.full(self.prompt_length, -100, dtype=np.int64), tgt))
-        inp = np.concatenate(
-            (
-                np.arange(self.prompt_length, dtype=np.int64) + self.prompt_length * global_task,
-                ctx,
-            )
-        )
-        return inp, tgt, inp.shape[0], context_inp, position_inp, segment_inp, task_inp
-
-    def __iter__(self):
-        while True:
-            ctx = self.ctx.read()
-            (
-                th_ctx,
-                th_tgt,
-                len_ctx,
-                context_ctx,
-                position_ctx,
-                segment_ctx,
-                task_ctx,
-            ) = self.__get_item_data(ctx)
-            yield th_ctx, th_tgt, len_ctx, context_ctx, position_ctx, segment_ctx, task_ctx
diff --git a/examples/README.md b/examples/README.md
new file mode 100644
index 0000000..8ecb69a
--- /dev/null
+++ b/examples/README.md
@@ -0,0 +1,16 @@
+# BMCook Examples
+
+## How to run examples
+
+1. enter examples path
+```shell
+$ git clone git@github.com:OpenBMB/BMCook.git
+$ cd BMCook/examples
+```
+
+2. run the scripts
+```shell
+$ bash cpm_live/cpm_ant_test.sh
+$ bash gpt/gpt2_test.sh
+$ bash t5/t5_test.sh
+```
\ No newline at end of file
diff --git a/cpm_live_example/README.md b/examples/cpm_live/README.md
similarity index 83%
rename from cpm_live_example/README.md
rename to examples/cpm_live/README.md
index ff8c988..3c200ae 100644
--- a/cpm_live_example/README.md
+++ b/examples/cpm_live/README.md
@@ -1,8 +1,8 @@
 # BMCook SPrune Guide for CPM-Ant Compression
 ## Overview
-BMCook Sprune is a main tool to compress CPM-Ant. The compression process is task-agnostic, in pre-training stage and shares the pre-training dataset. If you want to compress CPM-Ant for a specific downstream task, just prepare your task dataset in the way as [CPM-Ant](https://github.com/OpenBMB/CPM-Live/tree/master/cpm-live#4-feed-your-data) implements and follow the same compression process.
+BMCook Sprune is a main tool to compress CPM-Ant. 
 
-BMCook Sprune now supports two pruning granularities: layer level and head/dim level. Generally, you can set the number of layers you want for layer level pruning, then set the num_heads or dim_ff to do head/dim level pruning.
+It now supports two pruning granularities: layer level and head/dim level. Generally, you can set the number of layers you want for layer level pruning, then set the num_heads or dim_ff to do head/dim level pruning.
 
 More general pruning toolkit will be released soon!
 
@@ -64,7 +64,7 @@ The "train_mask" decides:
 
 The "prune_mode" decides:
 - att: train the head mask.
-- ffn: train the dim_ff mask. 
+- ffn: train the fim_ff mask. 
 
 If you train head mask first, you should provides the fixed mask in "heads_mask" when training fixed mask and vice versa. 
 
@@ -119,4 +119,4 @@ Scheme: num_heads: 16 -> 4, dim_ff: 700 -> 350
 5. do actual prune:
 ```shell
 python do_prune.py
-```
+```
\ No newline at end of file
diff --git a/examples/cpm_live/__init__.py b/examples/cpm_live/__init__.py
new file mode 100644
index 0000000..ec514b0
--- /dev/null
+++ b/examples/cpm_live/__init__.py
@@ -0,0 +1,4 @@
+from . import models
+from . import data
+from . import utils
+from . import tokenizers
\ No newline at end of file
diff --git a/examples/cpm_live/arguments.py b/examples/cpm_live/arguments.py
new file mode 100644
index 0000000..780319c
--- /dev/null
+++ b/examples/cpm_live/arguments.py
@@ -0,0 +1,31 @@
+# coding=utf-8
+# Copyright 2020 The OpenBMB team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model-config", type=str, help="model configuration file")
+    p.add_argument("--load", type=str, help="model checkpoint")
+    p.add_argument('--teacher-config', type=str, help="teacher configuration file")
+    p.add_argument('--load-teacher', type=str, help="model checkpoint")
+    p.add_argument('--log-interval', type=int, default=1, help='Log loss, LR, scale every n steps')
+    p.add_argument('--save-interval', type=int, default=1000, help='Interval steps between checkpoints')
+    p.add_argument('--start-lr', type=float, default=1e-4, help='Start learning rate')
+    p.add_argument('--data-path', default='openwebtext_document', help='Path to dataset')
+    p.add_argument('--cook-config', type=str, help='Path to BMCook config file')
+    
+    return p.parse_args()
+
diff --git a/examples/cpm_live/configs/combine.json b/examples/cpm_live/configs/combine.json
new file mode 100644
index 0000000..51348d4
--- /dev/null
+++ b/examples/cpm_live/configs/combine.json
@@ -0,0 +1,33 @@
+{
+    "distillation": {
+        "ce_scale": 0,
+
+        "mse_hidn_scale": 1,
+        "mse_hidn_module": ["[post]encoder.output_layernorm:[post]encoder.output_layernorm"],
+        "mse_hidn_proj": false
+    },
+
+    "pruning": {
+        "is_pruning": true,
+        "pruning_mask_path": "prune_example.bin",
+        "pruned_module": ["input_embedding"],
+        "mask_method": "sprune",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
+    },
+
+    "quantization": {
+        "is_quant": true
+    },
+
+    "MoEfication": {
+        "is_moefy": false,
+        "first_FFN_module": ["ffn.layernorm_before_ffn"]
+    }
+}
\ No newline at end of file
diff --git a/examples/cpm_live/configs/int8-kd.json b/examples/cpm_live/configs/int8-kd.json
new file mode 100644
index 0000000..51dad43
--- /dev/null
+++ b/examples/cpm_live/configs/int8-kd.json
@@ -0,0 +1,33 @@
+{
+    "distillation": {
+        "ce_scale": 0,
+
+        "mse_hidn_scale": 1,
+        "mse_hidn_module": ["[post]encoder.output_layernorm:[post]encoder.output_layernorm"],
+        "mse_hidn_proj": false
+    },
+
+    "pruning": {
+        "is_pruning": false,
+        "pruning_mask_path": "prune_example.bin",
+        "pruned_module": ["input_embedding"],
+        "mask_method": "sprune",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
+    },
+
+    "quantization": {
+        "is_quant": true
+    },
+
+    "MoEfication": {
+        "is_moefy": false,
+        "first_FFN_module": ["ffn.layernorm_before_ffn"]
+    }
+}
\ No newline at end of file
diff --git a/examples/cpm_live/configs/int8.json b/examples/cpm_live/configs/int8.json
new file mode 100644
index 0000000..fc2e642
--- /dev/null
+++ b/examples/cpm_live/configs/int8.json
@@ -0,0 +1,33 @@
+{
+    "distillation": {
+        "ce_scale": 0,
+
+        "mse_hidn_scale": 0,
+        "mse_hidn_module": ["[post]encoder.output_layernorm:[post]encoder.output_layernorm", "[pre]encoder.layers.9.ffn.layernorm_before_ffn:[pre]encoder.layers.9.ffn.layernorm_before_ffn"],
+        "mse_hidn_proj": false
+    },
+
+    "pruning": {
+        "is_pruning": false,
+        "pruning_mask_path": "prune_example.bin",
+        "pruned_module": ["input_embedding"],
+        "mask_method": "sprune",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
+    },
+
+    "quantization": {
+        "is_quant": true
+    },
+
+    "MoEfication": {
+        "is_moefy": false,
+        "first_FFN_module": ["ffn.layernorm_before_ffn"]
+    }
+}
diff --git a/examples/cpm_live/configs/moe.json b/examples/cpm_live/configs/moe.json
new file mode 100644
index 0000000..c618078
--- /dev/null
+++ b/examples/cpm_live/configs/moe.json
@@ -0,0 +1,33 @@
+{
+    "distillation": {
+        "ce_scale": 0,
+
+        "mse_hidn_scale": 0,
+        "mse_hidn_module": ["[post]encoder.output_layernorm:[post]encoder.output_layernorm", "[pre]encoder.layers.9.ffn.layernorm_before_ffn:[pre]encoder.layers.9.ffn.layernorm_before_ffn"],
+        "mse_hidn_proj": false
+    },
+
+    "pruning": {
+        "is_pruning": false,
+        "pruning_mask_path": "prune_example.bin",
+        "pruned_module": ["input_embedding"],
+        "mask_method": "sprune",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
+    },
+
+    "quantization": {
+        "is_quant": false
+    },
+
+    "MoEfication": {
+        "is_moefy": true,
+        "first_FFN_module": ["ffn.layernorm_before_ffn"]
+    }
+}
\ No newline at end of file
diff --git a/cpm_live_example/config/bmcook.json b/examples/cpm_live/configs/prune.json
similarity index 52%
rename from cpm_live_example/config/bmcook.json
rename to examples/cpm_live/configs/prune.json
index 5d91789..ddd18c8 100644
--- a/cpm_live_example/config/bmcook.json
+++ b/examples/cpm_live/configs/prune.json
@@ -1,6 +1,7 @@
 {
     "distillation": {
         "ce_scale": 0,
+
         "mse_hidn_scale": 0,
         "mse_hidn_module": ["[post]encoder.output_layernorm:[post]encoder.output_layernorm", "[pre]encoder.layers.9.ffn.layernorm_before_ffn:[pre]encoder.layers.9.ffn.layernorm_before_ffn"],
         "mse_hidn_proj": false
@@ -8,14 +9,17 @@
 
     "pruning": {
         "is_pruning": true,
-        "pruning_mask_path": "/yinxr/gongbt/BMCook/example/t5-mask-m4n2.bin",
-        "pruned_module": ["ffn.ffn.w_in.w.weight", 
-                            "ffn.ffn.w_out.weight", 
-                            "self_attention.project_q.weight",
-                            "self_attention.project_k.weight", 
-                            "self_attention.project_v.weight", 
-                            "self_attention.attention_out.weight"],
-        "mask_method": "fine-grained"
+        "pruning_mask_path": "prune_example.bin",
+        "pruned_module": ["input_embedding"],
+        "mask_method": "m4n2_1d",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
     },
 
     "quantization": {
diff --git a/examples/cpm_live/data/__init__.py b/examples/cpm_live/data/__init__.py
new file mode 100644
index 0000000..4e7dafa
--- /dev/null
+++ b/examples/cpm_live/data/__init__.py
@@ -0,0 +1,1 @@
+from .ant import CPMAnt_Dataset, DistributedMMapIndexedDataset, BatchPacker
\ No newline at end of file
diff --git a/examples/cpm_live/data/ant.py b/examples/cpm_live/data/ant.py
new file mode 100644
index 0000000..8685c52
--- /dev/null
+++ b/examples/cpm_live/data/ant.py
@@ -0,0 +1,340 @@
+import os
+import torch
+import random
+import struct
+import numpy as np
+import bmtrain as bmt
+import torch.utils.data as data
+from itertools import accumulate
+
+dtypes = {
+    1: np.uint8,
+    2: np.int8,
+    3: np.int16,
+    4: np.int32,
+    5: np.int64,
+    6: np.float32,
+    7: np.double,
+    8: np.uint16
+}
+
+def code(dtype):
+    for k in dtypes.keys():
+        if dtypes[k] == dtype:
+            return k
+    raise ValueError(dtype)
+
+def index_file_path(prefix_path):
+    return prefix_path + '.idx'
+
+def data_file_path(prefix_path):
+    return prefix_path + '.bin'
+
+class DistributedMMapIndexedDataset(torch.utils.data.Dataset):
+    class Index(object):
+        _HDR_MAGIC = b'MMIDIDX\x00\x00'
+        def __init__(self, path):
+            with open(path, 'rb') as stream:
+                magic_test = stream.read(9)
+                assert self._HDR_MAGIC == magic_test, (
+                    'Index file doesn\'t match expected format. '
+                    'Make sure that --dataset-impl is configured properly.'
+                )
+                version = struct.unpack('<Q', stream.read(8))
+                assert (1,) == version
+
+                dtype_code, = struct.unpack('<B', stream.read(1))
+                self._dtype = dtypes[dtype_code]
+                self._dtype_size = self._dtype().itemsize
+
+                self._len = struct.unpack('<Q', stream.read(8))[0]
+                self._doc_count = struct.unpack('<Q', stream.read(8))[0]
+                offset = stream.tell()
+
+            self._bin_buffer_mmap = np.memmap(path, mode='r', order='C')
+            self._bin_buffer = memoryview(self._bin_buffer_mmap)
+            self._sizes = np.frombuffer(
+                self._bin_buffer,
+                dtype=np.int32,
+                count=self._len,
+                offset=offset)
+            self._pointers = np.frombuffer(self._bin_buffer, dtype=np.int64, count=self._len,
+                                           offset=offset + self._sizes.nbytes)
+            self._doc_idx = np.frombuffer(self._bin_buffer, dtype=np.int64, count=self._doc_count,
+                                          offset=offset + self._sizes.nbytes + self._pointers.nbytes)
+
+        def __del__(self):
+            self._bin_buffer_mmap._mmap.close()
+            del self._bin_buffer_mmap
+
+        @property
+        def dtype(self):
+            return self._dtype
+
+        @property
+        def sizes(self):
+            return self._sizes
+
+        @property
+        def doc_idx(self):
+            return self._doc_idx
+
+        def __getitem__(self, i):
+            return self._pointers[i], self._sizes[i]
+
+        def __len__(self):
+            return self._len
+
+    def __init__(self, path, name, rank_number, rank_total, cache = None):
+        
+        super().__init__()
+
+        self._path = path
+        self._name = name
+        self._state = 0
+        if cache is not None:
+            self._cache = cache
+            os.makedirs(self._cache, exist_ok=True)
+        else:
+            self._cache = None
+        self._rank_total = rank_total
+        self._rank_number = rank_number
+        self._index = None
+        self._bin_buffer = None
+        self._bin_buffer_mmap = None
+        self.history = {-1:0}
+
+        self._do_init(self._path, self._name, self._cache, self._state)
+
+    def __getstate__(self):
+        return self._path + self._name + "_%d"%(self._state)
+
+    def __setstate__(self, state):
+        self._state = state
+        self._do_init(self._path, self._name, self._cache, self._state)
+
+    def _do_init(self, path, name, cache, state):
+        if self._bin_buffer_mmap is not None:
+            self._bin_buffer_mmap._mmap.close()
+            del self._bin_buffer_mmap
+        if self._index is not None:
+            del self._index
+
+        self._state = state
+
+        source_file = path + name + "_%d"%(self._state)
+
+        self._index = self.Index(index_file_path(source_file))
+        self._bin_buffer_mmap = np.memmap(data_file_path(source_file), mode='r', order='C')
+        self._bin_buffer = memoryview(self._bin_buffer_mmap)
+
+        self.history[state] = self.history[state - 1] + len(self._index) // self._rank_total
+        self.start = (len(self._index) // self._rank_total) * self._rank_number
+        print (self.history[state],"======================", self.start, self.history[self._state - 1])
+
+    def __del__(self):
+        if self._bin_buffer_mmap is not None:
+            self._bin_buffer_mmap._mmap.close()
+            del self._bin_buffer_mmap
+        if self._index is not None:
+            del self._index
+
+    def _next_file(self):
+        self._state += 1
+        bmt.print_rank(f"next_file: {self._state}")
+        self._do_init(self._path, self._name, self._cache, self._state)
+    
+    def __relative_idx(self, idx):
+        res = self.start + idx - self.history[self._state - 1]
+        return res
+
+    def __slice_item(self, start, stop):
+        ptr = self._index._pointers[self.__relative_idx(start)]
+        sizes = self._index._sizes[self.__relative_idx(start):self.__relative_idx(stop)]
+        offsets = list(accumulate(sizes))
+        np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype, count=sum(sizes), offset=ptr)
+        return np.split(np_array, offsets[:-1])
+
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            while idx >= self.history[self._state]:
+                self._next_file()
+            ptr, size = self._index[self.__relative_idx(idx)]
+            return np.frombuffer(self._bin_buffer, dtype=self._index.dtype, count=size, offset=ptr)
+        elif isinstance(idx, slice):
+            start, stop, step = idx.indices(2147483647)
+            assert step == 1 or step is None, "Slices into indexed_dataset must be contiguous"
+            if stop >= self.history[self._state]:
+                res_1 = self.__slice_item(start, self.history[self._state]) 
+                self._next_file()
+                res_2 = self.__slice_item(self.history[self._state - 1], stop)
+                return res_1 + res_2
+            else:
+                return self.__slice_item(start, stop)
+
+    @property
+    def sizes(self):
+        return self._index.sizes
+        
+    def exists(path):
+        return (os.path.exists(index_file_path(path)) and os.path.exists(data_file_path(path)))
+
+class CPMAnt_Dataset(data.Dataset):
+    def __init__(self, ctx, max_length = 1024, prompt_length = 32, tokenizer = None):
+        self.ctx = ctx
+        self.max_length = max_length + prompt_length
+        self.prompt_length = prompt_length
+        self.tokenizer = tokenizer
+
+    def __len__(self):
+        return len(self.ctx)
+
+    def __get_item_data(self, raw_data, index):
+
+        global_task = raw_data[0]
+        n_segment = raw_data[1]
+        len_info = n_segment * 3 + 2
+        segment_len  = raw_data[2:len_info:3] 
+        segment_type = raw_data[3:len_info:3]
+        segment_task = raw_data[4:len_info:3]
+        ctx = raw_data[len_info:]
+
+        if ctx.shape[0] > self.max_length - self.prompt_length:
+            return None, None, None, None, None, None, None
+        len_ctx = min(ctx.shape[0], self.max_length - self.prompt_length)
+
+        context_inp = np.full(len_ctx, True)
+        position_inp = np.arange(len_ctx, dtype=np.int64)
+        segment_inp = np.full(len_ctx, 0, dtype=np.int64)
+        task_inp = np.full(len_ctx, 0, dtype=np.int64)
+        tgt = np.full(len_ctx, -100, dtype=np.int64)
+
+        # for each segment
+        segment_begin = 0
+        for i in range(n_segment):
+            segment_end = segment_begin + segment_len[i]
+            task = segment_task[i]
+            # generate target
+            if task == 0:
+                num_mask = random.randint(1, segment_len[i] - 1)
+                mask_idx = np.random.choice(segment_len[i] - 1, num_mask, replace=False) + segment_begin
+                context_inp[mask_idx + 1] = False
+                task_inp[segment_begin:segment_end] = task
+                assert segment_type[i] == 1
+            elif task == 1:
+                num_mask = random.randint(1, segment_len[i] - 1)
+                context_inp[segment_end-num_mask:segment_end] = False
+                task_inp[segment_begin:segment_end] = task
+                assert segment_type[i] == 2
+            segment_inp[segment_begin:segment_end] = segment_type[i]
+            tgt[segment_begin : segment_end - 1] = np.where(
+                context_inp[segment_begin + 1 : segment_end],
+                -100,
+                ctx[segment_begin + 1 : segment_end]
+            )
+            segment_begin = segment_end
+        # prepend prompt segment
+        context_inp = np.concatenate((np.full(self.prompt_length, True), context_inp))
+        position_inp = np.concatenate((np.arange(self.prompt_length, dtype=np.int64), position_inp + self.prompt_length))
+        segment_inp = np.concatenate((np.full(self.prompt_length, 0, dtype=np.int64), segment_inp))
+        task_inp = np.concatenate((np.full(self.prompt_length, 0, dtype=np.int64), task_inp))
+        tgt = np.concatenate((np.full(self.prompt_length, -100, dtype=np.int64), tgt))
+        inp = np.concatenate((np.arange(self.prompt_length, dtype=np.int64) + self.prompt_length * global_task, ctx))
+        return inp, tgt, inp.shape[0], context_inp, position_inp, segment_inp, task_inp
+
+    def __getitem__(self, index):
+        ctx = self.ctx[index]
+        th_ctx, th_tgt, len_ctx, context_ctx, position_ctx, segment_ctx, task_ctx = \
+                self.__get_item_data(ctx, index)
+        return th_ctx, th_tgt, len_ctx, context_ctx, position_ctx, segment_ctx, task_ctx
+
+class BatchPacker:
+    def __init__(self, dataset, max_length, batch_size):
+        self.last_idx = 0
+        self.dataset = dataset
+        self.max_length = max_length
+        self.batch_size = batch_size
+    
+    def state(self):
+        return self.last_idx
+    
+    def load_state(self, state):
+        self.last_idx = state
+
+    def __iter__(self):
+        st = self.last_idx
+        ctx = []
+        tgt = []
+        context = []
+        position = []
+        segment = []
+        span = []
+        task_info = []
+
+        while True:
+            ctx_data, tgt_data, _len, context_data, position_data, segment_data, task_data = self.dataset[st]
+            st += 1
+            if ctx_data is None:
+                continue
+            assert _len <= self.max_length
+
+            ctx_data = ctx_data.astype("int64")
+            tgt_data = tgt_data.astype("int64")
+
+            for index in range(len(ctx)):
+                if span[index][-1] + _len < self.max_length:
+                    ctx[index][span[index][-1]:span[index][-1] + _len] = torch.from_numpy(ctx_data)[:_len].long()
+                    tgt[index][span[index][-1]:span[index][-1] + _len]= torch.from_numpy(tgt_data)[:_len].long()
+                    context[index][span[index][-1]:span[index][-1] + _len] = torch.from_numpy(context_data)[:_len].bool()
+                    position[index][span[index][-1]:span[index][-1] + _len] = torch.from_numpy(position_data)[:_len].long()
+                    segment[index][span[index][-1]:span[index][-1] + _len] = torch.from_numpy(segment_data)[:_len].long()
+                    task_info[index][span[index][-1]:span[index][-1] + _len] = torch.from_numpy(task_data)[:_len].long()
+                    span[index].append(span[index][-1] + _len)
+                    break
+            else:
+                _ctx = torch.zeros((self.max_length,), dtype=torch.long)
+                _ctx[:_len] = torch.from_numpy(ctx_data)[:_len].long()
+                _tgt = torch.full((self.max_length,), -100, dtype=torch.long)
+                _tgt[:_len] = torch.from_numpy(tgt_data)[:_len].long()
+                _context = torch.full((self.max_length,), False, dtype=torch.bool)
+                _context[:_len] = torch.from_numpy(context_data)[:_len].bool()
+                _position = torch.full((self.max_length,), False, dtype=torch.long)
+                _position[:_len] = torch.from_numpy(position_data)[:_len].long()
+                _segment = torch.full((self.max_length,), False, dtype=torch.long)
+                _segment[:_len] = torch.from_numpy(segment_data)[:_len].long()
+                _task_info = torch.full((self.max_length,), -1, dtype=torch.long)
+                _task_info[:_len] = torch.from_numpy(task_data)[:_len].long()
+                ctx.append(_ctx)
+                tgt.append(_tgt)
+                context.append(_context)
+                position.append(_position)
+                segment.append(_segment)
+                task_info.append(_task_info)
+                span.append([_len])
+
+            if len(ctx) > self.batch_size:
+                _span = torch.zeros((self.batch_size, self.max_length + 1), dtype=torch.long)
+                for bindex in range(self.batch_size):
+                    for sindex in span[bindex]:
+                        _span[bindex][sindex] = 1
+                
+                self.last_idx = st
+                yield {
+                    "ctx": torch.stack(ctx[:self.batch_size]),
+                    "tgt": torch.stack(tgt[:self.batch_size]),
+                    "context": torch.stack(context[:self.batch_size]),
+                    "segment": torch.stack(segment[:self.batch_size]),
+                    "position": torch.stack(position[:self.batch_size]),
+                    "span": torch.cumsum(_span, dim=-1)[:,:-1],
+                    "len_ctx": torch.LongTensor([it[-1] for it in span[:self.batch_size]]),
+                    "task": torch.stack(task_info[:self.batch_size]),
+                }
+
+                ctx = ctx[self.batch_size:]
+                tgt = tgt[self.batch_size:]
+                context = context[self.batch_size:]
+                segment = segment[self.batch_size:]
+                position = position[self.batch_size:]
+                span = span[self.batch_size:]
+                task_info = task_info[self.batch_size:]
+
diff --git a/examples/cpm_live/layers/__init__.py b/examples/cpm_live/layers/__init__.py
new file mode 100644
index 0000000..4ff22c1
--- /dev/null
+++ b/examples/cpm_live/layers/__init__.py
@@ -0,0 +1,8 @@
+from .embedding import Embedding
+from .position_embedding import SegmentPositionEmbedding
+from .linear import Linear
+from .layernorm import LayerNorm
+from .attention import Attention
+from .feedforward import FeedForward
+from .blocks import TransformerBlock
+from .transformer import Encoder
diff --git a/cpm_live_example/cpm_live/layers/prune_attention.py b/examples/cpm_live/layers/attention.py
similarity index 93%
rename from cpm_live_example/cpm_live/layers/prune_attention.py
rename to examples/cpm_live/layers/attention.py
index 9deec78..c227037 100644
--- a/cpm_live_example/cpm_live/layers/prune_attention.py
+++ b/examples/cpm_live/layers/attention.py
@@ -57,7 +57,6 @@ def forward(
         position_bias: torch.Tensor,
         use_cache: bool = False,
         past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        att_z = None,
     ):
         """This model inherits from bmt.DistributedModule.
         Args:
@@ -81,12 +80,6 @@ def forward(
         h_k = h_k.view(batch_size, len_k, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
         h_v = h_v.view(batch_size, len_k, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
 
-        #prune
-        if att_z is not None:
-            h_q = h_q * att_z.view(self.num_heads, 1, 1)
-            h_k = h_k * att_z.view(self.num_heads, 1, 1)
-            h_v = h_v * att_z.view(self.num_heads, 1, 1)
-
         if past_kv is not None:
             h_k = torch.cat([past_kv[0], h_k], dim=-2)
             h_v = torch.cat([past_kv[1], h_v], dim=-2)
@@ -116,9 +109,6 @@ def forward(
         score = torch.matmul(score, h_v)
 
         score = score.view(batch_size, self.num_heads, len_q, self.dim_head).permute(0, 2, 1, 3)
-        # prune
-        if att_z is not None:
-            score = score * att_z.view(self.num_heads, 1)
         score = score.contiguous().view(batch_size, len_q, self.num_heads * self.dim_head)
 
         score = self.attention_out(score)
diff --git a/cpm_live_example/cpm_live/layers/prune_blocks.py b/examples/cpm_live/layers/blocks.py
similarity index 78%
rename from cpm_live_example/cpm_live/layers/prune_blocks.py
rename to examples/cpm_live/layers/blocks.py
index efbd4c3..b682541 100644
--- a/cpm_live_example/cpm_live/layers/prune_blocks.py
+++ b/examples/cpm_live/layers/blocks.py
@@ -1,9 +1,24 @@
+# coding=utf-8
+# Copyright 2022 The OpenBMB team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Optional, Tuple
 import torch
 import bmtrain as bmt
 from .layernorm import LayerNorm
-from .prune_attention import Attention
-from .prune_feedforward import FeedForward
+from .attention import Attention
+from .feedforward import FeedForward
 
 
 class SelfAttentionBlock(bmt.DistributedModule):
@@ -56,8 +71,6 @@ def forward(
         position_bias: Optional[torch.Tensor] = None,
         use_cache: bool = False,
         past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        layer_z: torch.Tensor = None,
-        att_z: torch.Tensor = None,
     ):
         """
         Args:
@@ -70,7 +83,7 @@ def forward(
 
         """  # noqa: E501
         x = self.layernorm_before_attention(hidden_states)
-        x = self.self_attention(x, x, attention_mask, position_bias, use_cache, past_key_value, att_z=att_z)
+        x = self.self_attention(x, x, attention_mask, position_bias, use_cache, past_key_value)
         if use_cache:
             x, current_key_value = x
         else:
@@ -78,8 +91,6 @@ def forward(
 
         if self.dropout is not None:
             x = self.dropout(x)
-        if layer_z is not None:
-            x = x * layer_z
         hidden_states = hidden_states + x
 
         if use_cache:
@@ -130,8 +141,6 @@ def __init__(
     def forward(
         self,
         hidden_states: torch.Tensor,
-        layer_z = None,
-        ffn_z = None,
     ):
         """
         Args:
@@ -142,16 +151,14 @@ def forward(
 
         """  # noqa: E501
         x = self.layernorm_before_ffn(hidden_states)
-        x = self.ffn(x, ffn_z=ffn_z)
+        x = self.ffn(x)
         if self.dropout is not None:
             x = self.dropout(x)
-        if layer_z is not None:
-            x = x * layer_z
         hidden_states = hidden_states + x
         return hidden_states
 
 
-class TransformerBlock(torch.nn.Module):
+class TransformerBlock(bmt.DistributedModule):
     """The whole transformer block. A sequence of operation. Consists of self-attention block[, cross-attention block] and feed-forward block.
 
     Args:
@@ -170,7 +177,6 @@ def __init__(
         dim_ff: int,
         num_heads: int,
         dim_head: int,
-        index: float,
         dtype=torch.half,
         eps: float = 1e-6,
         dropout_p: Optional[float] = None,
@@ -180,7 +186,6 @@ def __init__(
         super().__init__()
         self.mask_att = mask_att
         self.mask_ffn = mask_ffn
-        self.index = index
 
         if not self.mask_att:
             self.self_att = SelfAttentionBlock(
@@ -201,14 +206,6 @@ def __init__(
                 dropout_p=dropout_p,
             )
 
-        self.att_z = [1., 1., 1., 0., 1., 1., 1., 1., 1., 1., 0., 1., 1., 0., 0., 0., 1., 0.,
-        1., 0., 1., 0., 0., 1., 1., 0., 0., 1., 0., 1., 1., 1., 0., 1., 0., 0., 1., 0., 0., 
-        1., 1., 0., 1., 0., 0., 0., 0., 0.]
-        self.ffn_z = [1., 0., 0., 0., 0., 1., 0., 1., 0., 1., 0., 1., 0., 0., 0., 0., 1., 1., 
-        1., 0., 1., 0., 1., 0., 1., 0., 0., 1., 1., 0., 0., 1., 0., 0., 0., 0., 1., 1., 0., 
-        1., 1., 0., 0., 0., 1., 1., 0., 1.]
-        
-
     def forward(
         self,
         self_hidden_states: torch.Tensor,
@@ -216,7 +213,6 @@ def forward(
         self_position_bias: Optional[torch.Tensor] = None,
         use_cache: bool = False,
         past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        masks = None,
     ):
         """
         Args:
@@ -230,24 +226,13 @@ def forward(
         """  # noqa: E501
         # (batch, dim_model, seq_self)
         current_key_value = None
-        att_z, ffn_z, att_layer_z, ffn_layer_z = None, None, None, None
         if not self.mask_att:
-            if masks is not None:
-                att_index = int(sum(self.att_z[:self.index]))
-                if 'self_att_layer_z' in masks:
-                    att_layer_z = masks['self_att_layer_z'][att_index].clone().detach()
-                    att_layer_z = att_layer_z.to(dtype=torch.half, device="cuda")
-                if 'heads_z' in masks:
-                    att_z = masks['heads_z'][att_index].clone().detach()
-                    att_z = att_z.to(dtype=torch.half, device="cuda")
             hidden_states = self.self_att(
                 self_hidden_states,
                 attention_mask=self_attention_mask,
                 position_bias=self_position_bias,
                 use_cache=use_cache,
                 past_key_value=past_key_value,
-                layer_z=att_layer_z,
-                att_z=att_z
             )
             if use_cache:
                 hidden_states, current_key_value = hidden_states
@@ -256,15 +241,7 @@ def forward(
 
         # (batch, dim_model, seq_self)
         if not self.mask_ffn:
-            if masks is not None:
-                ffn_index = int(sum(self.ffn_z[:self.index]))
-                if 'ffn_layer_z' in masks:
-                    ffn_layer_z = masks['ffn_layer_z'][ffn_index].clone().detach()
-                    ffn_layer_z = ffn_layer_z.to(dtype=torch.half, device="cuda")
-                if 'dimff_z' in masks:
-                    ffn_z = masks['dimff_z'][ffn_index].clone().detach()
-                    ffn_z = ffn_z.to(dtype=torch.half, device="cuda")
-            hidden_states = self.ffn(hidden_states, layer_z=ffn_layer_z, ffn_z=ffn_z)
+            hidden_states = self.ffn(hidden_states)
 
         if use_cache:
             return hidden_states, current_key_value
diff --git a/cpm_live_example/cpm_live/layers/embedding.py b/examples/cpm_live/layers/embedding.py
similarity index 100%
rename from cpm_live_example/cpm_live/layers/embedding.py
rename to examples/cpm_live/layers/embedding.py
diff --git a/cpm_live_example/cpm_live/layers/prune_feedforward.py b/examples/cpm_live/layers/feedforward.py
similarity index 97%
rename from cpm_live_example/cpm_live/layers/prune_feedforward.py
rename to examples/cpm_live/layers/feedforward.py
index 6bd66e0..8dd3c23 100644
--- a/cpm_live_example/cpm_live/layers/prune_feedforward.py
+++ b/examples/cpm_live/layers/feedforward.py
@@ -104,7 +104,7 @@ def __init__(
             scale_before=True,
         )
 
-    def forward(self, x: torch.Tensor, ffn_z = None):
+    def forward(self, x: torch.Tensor):
         """
         Args:
             x (:obj:`torch.Tensor` of shape ``(batch, seq_len, dim_in)``): The input of feed-forward module.
@@ -114,9 +114,6 @@ def forward(self, x: torch.Tensor, ffn_z = None):
         """  # noqa: E501
         x = self.w_in(x)
 
-        if ffn_z is not None:
-            x = x * ffn_z
-
         if self.dropout is not None:
             x = self.dropout(x)
 
diff --git a/cpm_live_example/cpm_live/layers/layernorm.py b/examples/cpm_live/layers/layernorm.py
similarity index 100%
rename from cpm_live_example/cpm_live/layers/layernorm.py
rename to examples/cpm_live/layers/layernorm.py
diff --git a/cpm_live_example/cpm_live/layers/linear.py b/examples/cpm_live/layers/linear.py
similarity index 100%
rename from cpm_live_example/cpm_live/layers/linear.py
rename to examples/cpm_live/layers/linear.py
diff --git a/cpm_live_example/cpm_live/layers/position_embedding.py b/examples/cpm_live/layers/position_embedding.py
similarity index 100%
rename from cpm_live_example/cpm_live/layers/position_embedding.py
rename to examples/cpm_live/layers/position_embedding.py
diff --git a/cpm_live_example/cpm_live/layers/prune_transformer.py b/examples/cpm_live/layers/transformer.py
similarity index 86%
rename from cpm_live_example/cpm_live/layers/prune_transformer.py
rename to examples/cpm_live/layers/transformer.py
index 6c79461..fef5116 100644
--- a/cpm_live_example/cpm_live/layers/prune_transformer.py
+++ b/examples/cpm_live/layers/transformer.py
@@ -1,8 +1,23 @@
+# coding=utf-8
+# Copyright 2022 The OpenBMB team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import torch
 import bmtrain as bmt
 from typing import Optional, List, Tuple
 
-from .prune_blocks import TransformerBlock
+from .blocks import TransformerBlock
 from .layernorm import LayerNorm
 
 
@@ -56,7 +71,6 @@ def __init__(
                         dim_ff=dim_ff,
                         num_heads=num_heads,
                         dim_head=dim_head,
-                        index=ith,
                         dtype=dtype,
                         eps=eps,
                         dropout_p=dropout_p,
@@ -77,7 +91,6 @@ def forward(
         position_bias: torch.Tensor,
         use_cache: bool = False,
         past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
-        masks = None,
     ):
         """
         Args:
@@ -90,7 +103,7 @@ def forward(
 
         """  # noqa: E501
         if not use_cache:
-            hidden_states = self.layers(hidden_states, attention_mask, position_bias, False, None, masks)
+            hidden_states = self.layers(hidden_states, attention_mask, position_bias)
             hidden_states = self.output_layernorm(hidden_states)
             return hidden_states
         else:
@@ -103,7 +116,6 @@ def forward(
                         position_bias,
                         past_key_value=past_key_values[i] if past_key_values else None,
                         use_cache=use_cache,
-                        masks=masks,
                     )
                     if use_cache:
                         current_key_values.append(hidden_states[1])
diff --git a/cpm_live_example/cpm_live/models/__init__.py b/examples/cpm_live/models/__init__.py
similarity index 100%
rename from cpm_live_example/cpm_live/models/__init__.py
rename to examples/cpm_live/models/__init__.py
diff --git a/cpm_live_example/cpm_live/models/ant.py b/examples/cpm_live/models/ant.py
similarity index 99%
rename from cpm_live_example/cpm_live/models/ant.py
rename to examples/cpm_live/models/ant.py
index 3d718ff..ac93538 100644
--- a/cpm_live_example/cpm_live/models/ant.py
+++ b/examples/cpm_live/models/ant.py
@@ -52,6 +52,8 @@ def __init__(self, config: CPMAntConfig):
 
         super().__init__()
 
+        self.config = config
+
         self.encoder = Encoder(
             num_layers=config.num_layers,
             dim_model=config.dim_model,
@@ -134,7 +136,7 @@ def forward(
             )
 
         position_bias = self.position_bias(position, position, segment, segment)
-        hidden_states = self.encoder(hidden_states, attention_mask, position_bias, masks=masks)
+        hidden_states = self.encoder(hidden_states, attention_mask, position_bias)
 
         logits = self.input_embedding.projection(hidden_states)
         return logits, hidden_states
diff --git a/cpm_live_example/cpm_live/tokenizers/__init__.py b/examples/cpm_live/tokenizers/__init__.py
similarity index 100%
rename from cpm_live_example/cpm_live/tokenizers/__init__.py
rename to examples/cpm_live/tokenizers/__init__.py
diff --git a/cpm_live_example/cpm_live/tokenizers/ant.py b/examples/cpm_live/tokenizers/ant.py
similarity index 100%
rename from cpm_live_example/cpm_live/tokenizers/ant.py
rename to examples/cpm_live/tokenizers/ant.py
diff --git a/cpm_live_example/cpm_live/utils/__init__.py b/examples/cpm_live/utils/__init__.py
similarity index 100%
rename from cpm_live_example/cpm_live/utils/__init__.py
rename to examples/cpm_live/utils/__init__.py
diff --git a/cpm_live_example/cpm_live/utils/config.py b/examples/cpm_live/utils/config.py
similarity index 100%
rename from cpm_live_example/cpm_live/utils/config.py
rename to examples/cpm_live/utils/config.py
diff --git a/cpm_live_example/cpm_live/vocabs/ant.txt b/examples/cpm_live/vocabs/ant.txt
similarity index 100%
rename from cpm_live_example/cpm_live/vocabs/ant.txt
rename to examples/cpm_live/vocabs/ant.txt
diff --git a/gpt-example/README.md b/examples/gpt/README.md
similarity index 100%
rename from gpt-example/README.md
rename to examples/gpt/README.md
diff --git a/gpt-example/configs/gpt2-combine.json b/examples/gpt/configs/gpt2-combine.json
similarity index 61%
rename from gpt-example/configs/gpt2-combine.json
rename to examples/gpt/configs/gpt2-combine.json
index b1ed7a8..51348d4 100644
--- a/gpt-example/configs/gpt2-combine.json
+++ b/examples/gpt/configs/gpt2-combine.json
@@ -11,7 +11,15 @@
         "is_pruning": true,
         "pruning_mask_path": "prune_example.bin",
         "pruned_module": ["input_embedding"],
-        "mask_method": "m4n2_1d"
+        "mask_method": "sprune",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
     },
 
     "quantization": {
diff --git a/gpt-example/configs/gpt2-int8-kd.json b/examples/gpt/configs/gpt2-int8-kd.json
similarity index 55%
rename from gpt-example/configs/gpt2-int8-kd.json
rename to examples/gpt/configs/gpt2-int8-kd.json
index 94eccf9..51dad43 100644
--- a/gpt-example/configs/gpt2-int8-kd.json
+++ b/examples/gpt/configs/gpt2-int8-kd.json
@@ -10,8 +10,16 @@
     "pruning": {
         "is_pruning": false,
         "pruning_mask_path": "prune_example.bin",
-        "pruned_module": ["ffn.ffn.w_in.w.weight", "ffn.ffn.w_out.weight", "input_embedding"],
-        "mask_method": "m4n2_1d"
+        "pruned_module": ["input_embedding"],
+        "mask_method": "sprune",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
     },
 
     "quantization": {
diff --git a/gpt-example/configs/gpt2-int8.json b/examples/gpt/configs/gpt2-int8.json
similarity index 60%
rename from gpt-example/configs/gpt2-int8.json
rename to examples/gpt/configs/gpt2-int8.json
index 57fd3b7..fc2e642 100644
--- a/gpt-example/configs/gpt2-int8.json
+++ b/examples/gpt/configs/gpt2-int8.json
@@ -10,8 +10,16 @@
     "pruning": {
         "is_pruning": false,
         "pruning_mask_path": "prune_example.bin",
-        "pruned_module": ["ffn.ffn.w_in.w.weight", "ffn.ffn.w_out.weight", "input_embedding"],
-        "mask_method": "m4n2_1d"
+        "pruned_module": ["input_embedding"],
+        "mask_method": "sprune",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
     },
 
     "quantization": {
diff --git a/gpt-example/configs/gpt2-moe.json b/examples/gpt/configs/gpt2-moe.json
similarity index 55%
rename from gpt-example/configs/gpt2-moe.json
rename to examples/gpt/configs/gpt2-moe.json
index b8c3e35..c618078 100644
--- a/gpt-example/configs/gpt2-moe.json
+++ b/examples/gpt/configs/gpt2-moe.json
@@ -9,9 +9,17 @@
 
     "pruning": {
         "is_pruning": false,
-        "pruning_mask_path": "/home/zhangzhengyan//data/BMCook/example/prune_example.bin",
-        "pruned_module": ["ffn.ffn.w_in.w.weight", "ffn.ffn.w_out.weight", "input_embedding"],
-        "mask_method": "m4n2_1d"
+        "pruning_mask_path": "prune_example.bin",
+        "pruned_module": ["input_embedding"],
+        "mask_method": "sprune",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
     },
 
     "quantization": {
diff --git a/gpt-example/configs/gpt2-prune.json b/examples/gpt/configs/gpt2-prune.json
similarity index 100%
rename from gpt-example/configs/gpt2-prune.json
rename to examples/gpt/configs/gpt2-prune.json
diff --git a/gpt-example/data/__init__.py b/examples/gpt/data/__init__.py
similarity index 100%
rename from gpt-example/data/__init__.py
rename to examples/gpt/data/__init__.py
diff --git a/gpt-example/data/dataset.py b/examples/gpt/data/dataset.py
similarity index 100%
rename from gpt-example/data/dataset.py
rename to examples/gpt/data/dataset.py
diff --git a/gpt-example/data/indexed_dataset.py b/examples/gpt/data/indexed_dataset.py
similarity index 100%
rename from gpt-example/data/indexed_dataset.py
rename to examples/gpt/data/indexed_dataset.py
diff --git a/gpt-example/preprocess_data.py b/examples/gpt/preprocess_data.py
similarity index 100%
rename from gpt-example/preprocess_data.py
rename to examples/gpt/preprocess_data.py
diff --git a/examples/t5/README.md b/examples/t5/README.md
new file mode 100644
index 0000000..fd08f1e
--- /dev/null
+++ b/examples/t5/README.md
@@ -0,0 +1,7 @@
+You can generate the pre-training data as follows:
+
+```
+  python preprocess.py --input your_pretraining_data --output-prefix your_output_prefix
+```
+
+Each line of the input file contains one document.
diff --git a/examples/t5/configs/t5-combine.json b/examples/t5/configs/t5-combine.json
new file mode 100644
index 0000000..51348d4
--- /dev/null
+++ b/examples/t5/configs/t5-combine.json
@@ -0,0 +1,33 @@
+{
+    "distillation": {
+        "ce_scale": 0,
+
+        "mse_hidn_scale": 1,
+        "mse_hidn_module": ["[post]encoder.output_layernorm:[post]encoder.output_layernorm"],
+        "mse_hidn_proj": false
+    },
+
+    "pruning": {
+        "is_pruning": true,
+        "pruning_mask_path": "prune_example.bin",
+        "pruned_module": ["input_embedding"],
+        "mask_method": "sprune",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
+    },
+
+    "quantization": {
+        "is_quant": true
+    },
+
+    "MoEfication": {
+        "is_moefy": false,
+        "first_FFN_module": ["ffn.layernorm_before_ffn"]
+    }
+}
\ No newline at end of file
diff --git a/examples/t5/configs/t5-int8-kd.json b/examples/t5/configs/t5-int8-kd.json
new file mode 100644
index 0000000..51dad43
--- /dev/null
+++ b/examples/t5/configs/t5-int8-kd.json
@@ -0,0 +1,33 @@
+{
+    "distillation": {
+        "ce_scale": 0,
+
+        "mse_hidn_scale": 1,
+        "mse_hidn_module": ["[post]encoder.output_layernorm:[post]encoder.output_layernorm"],
+        "mse_hidn_proj": false
+    },
+
+    "pruning": {
+        "is_pruning": false,
+        "pruning_mask_path": "prune_example.bin",
+        "pruned_module": ["input_embedding"],
+        "mask_method": "sprune",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
+    },
+
+    "quantization": {
+        "is_quant": true
+    },
+
+    "MoEfication": {
+        "is_moefy": false,
+        "first_FFN_module": ["ffn.layernorm_before_ffn"]
+    }
+}
\ No newline at end of file
diff --git a/examples/t5/configs/t5-int8.json b/examples/t5/configs/t5-int8.json
new file mode 100644
index 0000000..fc2e642
--- /dev/null
+++ b/examples/t5/configs/t5-int8.json
@@ -0,0 +1,33 @@
+{
+    "distillation": {
+        "ce_scale": 0,
+
+        "mse_hidn_scale": 0,
+        "mse_hidn_module": ["[post]encoder.output_layernorm:[post]encoder.output_layernorm", "[pre]encoder.layers.9.ffn.layernorm_before_ffn:[pre]encoder.layers.9.ffn.layernorm_before_ffn"],
+        "mse_hidn_proj": false
+    },
+
+    "pruning": {
+        "is_pruning": false,
+        "pruning_mask_path": "prune_example.bin",
+        "pruned_module": ["input_embedding"],
+        "mask_method": "sprune",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
+    },
+
+    "quantization": {
+        "is_quant": true
+    },
+
+    "MoEfication": {
+        "is_moefy": false,
+        "first_FFN_module": ["ffn.layernorm_before_ffn"]
+    }
+}
diff --git a/examples/t5/configs/t5-moe.json b/examples/t5/configs/t5-moe.json
new file mode 100644
index 0000000..c618078
--- /dev/null
+++ b/examples/t5/configs/t5-moe.json
@@ -0,0 +1,33 @@
+{
+    "distillation": {
+        "ce_scale": 0,
+
+        "mse_hidn_scale": 0,
+        "mse_hidn_module": ["[post]encoder.output_layernorm:[post]encoder.output_layernorm", "[pre]encoder.layers.9.ffn.layernorm_before_ffn:[pre]encoder.layers.9.ffn.layernorm_before_ffn"],
+        "mse_hidn_proj": false
+    },
+
+    "pruning": {
+        "is_pruning": false,
+        "pruning_mask_path": "prune_example.bin",
+        "pruned_module": ["input_embedding"],
+        "mask_method": "sprune",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
+    },
+
+    "quantization": {
+        "is_quant": false
+    },
+
+    "MoEfication": {
+        "is_moefy": true,
+        "first_FFN_module": ["ffn.layernorm_before_ffn"]
+    }
+}
\ No newline at end of file
diff --git a/examples/t5/configs/t5-prune.json b/examples/t5/configs/t5-prune.json
new file mode 100644
index 0000000..ddd18c8
--- /dev/null
+++ b/examples/t5/configs/t5-prune.json
@@ -0,0 +1,33 @@
+{
+    "distillation": {
+        "ce_scale": 0,
+
+        "mse_hidn_scale": 0,
+        "mse_hidn_module": ["[post]encoder.output_layernorm:[post]encoder.output_layernorm", "[pre]encoder.layers.9.ffn.layernorm_before_ffn:[pre]encoder.layers.9.ffn.layernorm_before_ffn"],
+        "mse_hidn_proj": false
+    },
+
+    "pruning": {
+        "is_pruning": true,
+        "pruning_mask_path": "prune_example.bin",
+        "pruned_module": ["input_embedding"],
+        "mask_method": "m4n2_1d",
+        "sprune": {
+                "criterion": "l0",
+                "training_mask": ["transformer", "num_heads"],
+                "fixed_mask_path": "",
+                "mask_mode": "train_mask",
+                "target_mode": "sparsity",
+                "target_sparsity": 0.5
+                }
+    },
+
+    "quantization": {
+        "is_quant": false
+    },
+
+    "MoEfication": {
+        "is_moefy": false,
+        "first_FFN_module": ["ffn.layernorm_before_ffn"]
+    }
+}
\ No newline at end of file
diff --git a/examples/t5/data/__init__.py b/examples/t5/data/__init__.py
new file mode 100644
index 0000000..5cb59a2
--- /dev/null
+++ b/examples/t5/data/__init__.py
@@ -0,0 +1,2 @@
+from .indexed_dataset import MMapIndexedDataset
+from .dataset import Dataset
\ No newline at end of file
diff --git a/examples/t5/data/dataset.py b/examples/t5/data/dataset.py
new file mode 100644
index 0000000..49d7d32
--- /dev/null
+++ b/examples/t5/data/dataset.py
@@ -0,0 +1,41 @@
+import numpy as np
+import torch
+import torch.utils.data as data
+from . import MMapIndexedDataset
+
+
+class Dataset(data.Dataset):
+    def __init__(self, ctx_context : MMapIndexedDataset, ctx_target : MMapIndexedDataset, ctx_len, tar_len):
+        self.ctx_context = ctx_context
+        self.ctx_target = ctx_target
+        self.ctx_len = ctx_len
+        self.tar_len = tar_len
+
+    def __len__(self):
+        return len(self.ctx_context)
+    
+    def __getitem__(self, index):
+        ctx_context = self.ctx_context[index]
+        ctx_target = self.ctx_target[index]
+
+        if ctx_context.shape[0] > self.ctx_len or ctx_target.shape[0] > self.tar_len:
+            return None
+
+        len_ctx_context = ctx_context.shape[0]
+        len_ctx_target = ctx_target.shape[0]
+        
+        th_ctx_context = torch.zeros(self.ctx_len, dtype=torch.long)
+        th_ctx_context[:len_ctx_context] = torch.from_numpy(ctx_context.astype(np.int))[:len_ctx_context].long()
+
+        th_ctx_target = torch.zeros(self.tar_len, dtype=torch.long)
+        th_ctx_target[:len_ctx_target-1] = torch.from_numpy(ctx_target.astype(np.int))[:len_ctx_target-1].long()
+
+        th_target = torch.zeros(self.tar_len, dtype=torch.long) * -100
+        th_target[:len_ctx_target-1] = torch.from_numpy(ctx_target.astype(np.int))[1:len_ctx_target].long()
+        return {
+            "ctx_context": th_ctx_context,
+            "len_ctx_context": len_ctx_context,
+            "ctx_target": th_ctx_target,
+            "len_ctx_target": len_ctx_target-1,
+            'target': th_target
+        }
\ No newline at end of file
diff --git a/examples/t5/data/indexed_dataset.py b/examples/t5/data/indexed_dataset.py
new file mode 100644
index 0000000..d82354b
--- /dev/null
+++ b/examples/t5/data/indexed_dataset.py
@@ -0,0 +1,571 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+# copied from fairseq/fairseq/data/indexed_dataset.py
+# Removed IndexedRawTextDataset since it relied on Fairseq dictionary
+# other slight modifications to remove fairseq dependencies
+# Added document index to index file and made it accessible.
+#    An empty sentence no longer separates documents.
+
+from functools import lru_cache
+import os
+import shutil
+import struct
+from itertools import accumulate
+import bmtrain as bmt
+
+import numpy as np
+import torch
+
+
+def __best_fitting_dtype(vocab_size=None):
+    if vocab_size is not None and vocab_size < 65500:
+        return np.uint16
+    else:
+        return np.int32
+
+
+def get_available_dataset_impl():
+    return ['lazy', 'cached', 'mmap']
+
+
+def infer_dataset_impl(path):
+    if IndexedDataset.exists(path):
+        with open(index_file_path(path), 'rb') as f:
+            magic = f.read(8)
+            if magic == IndexedDataset._HDR_MAGIC:
+                return 'cached'
+            elif magic == MMapIndexedDataset.Index._HDR_MAGIC[:8]:
+                return 'mmap'
+            else:
+                return None
+    else:
+        bmt.print_rank(f"Dataset does not exist: {path}")
+        bmt.print_rank("Path should be a basename that both .idx and .bin can be appended to get full filenames.")
+        return None
+
+
+def make_builder(out_file, impl, dtype):
+    if impl == 'mmap':
+        return MMapIndexedDatasetBuilder(out_file, dtype=dtype)
+    else:
+        return IndexedDatasetBuilder(out_file)
+
+
+def make_dataset(path, impl, skip_warmup=False):
+    if not IndexedDataset.exists(path):
+        bmt.print_rank(f"Dataset does not exist: {path}")
+        bmt.print_rank("Path should be a basename that both .idx and .bin can be appended to get full filenames.")
+        return None
+    if impl == 'infer':
+        impl = infer_dataset_impl(path)
+    if impl == 'lazy' and IndexedDataset.exists(path):
+        return IndexedDataset(path)
+    elif impl == 'cached' and IndexedDataset.exists(path):
+        return IndexedCachedDataset(path)
+    elif impl == 'mmap' and MMapIndexedDataset.exists(path):
+        return MMapIndexedDataset(path, skip_warmup)
+    bmt.print_rank(f"Unknown dataset implementation: {impl}")
+    return None
+
+
+def dataset_exists(path, impl):
+    if impl == 'mmap':
+        return MMapIndexedDataset.exists(path)
+    else:
+        return IndexedDataset.exists(path)
+
+
+def read_longs(f, n):
+    a = np.empty(n, dtype=np.int64)
+    f.readinto(a)
+    return a
+
+
+def write_longs(f, a):
+    f.write(np.array(a, dtype=np.int64))
+
+
+dtypes = {
+    1: np.uint8,
+    2: np.int8,
+    3: np.int16,
+    4: np.int32,
+    5: np.int64,
+    6: np.float,
+    7: np.double,
+    8: np.uint16
+}
+
+
+def code(dtype):
+    for k in dtypes.keys():
+        if dtypes[k] == dtype:
+            return k
+    raise ValueError(dtype)
+
+
+def index_file_path(prefix_path):
+    return prefix_path + '.idx'
+
+
+def data_file_path(prefix_path):
+    return prefix_path + '.bin'
+
+
+def create_doc_idx(sizes):
+    doc_idx = [0]
+    for i, s in enumerate(sizes):
+        if s == 0:
+            doc_idx.append(i + 1)
+    return doc_idx
+
+
+class IndexedDataset(torch.utils.data.Dataset):
+    """Loader for IndexedDataset"""
+    _HDR_MAGIC = b'TNTIDX\x00\x00'
+
+    def __init__(self, path):
+        super().__init__()
+        self.path = path
+        self.data_file = None
+        self.read_index(path)
+
+    def read_index(self, path):
+        with open(index_file_path(path), 'rb') as f:
+            magic = f.read(8)
+            assert magic == self._HDR_MAGIC, (
+                'Index file doesn\'t match expected format. '
+                'Make sure that --dataset-impl is configured properly.'
+            )
+            version = f.read(8)
+            assert struct.unpack('<Q', version) == (1,)
+            code, self.element_size = struct.unpack('<QQ', f.read(16))
+            self.dtype = dtypes[code]
+            self._len, self.s = struct.unpack('<QQ', f.read(16))
+            self.doc_count = struct.unpack('<Q', f.read(8))
+            self.dim_offsets = read_longs(f, self._len + 1)
+            self.data_offsets = read_longs(f, self._len + 1)
+            self.sizes = read_longs(f, self.s)
+            self.doc_idx = read_longs(f, self.doc_count)
+
+    def read_data(self, path):
+        self.data_file = open(data_file_path(path), 'rb', buffering=0)
+
+    def check_index(self, i):
+        if i < 0 or i >= self._len:
+            raise IndexError('index out of range')
+
+    def __del__(self):
+        if self.data_file:
+            self.data_file.close()
+
+    # @lru_cache(maxsize=8)
+    def __getitem__(self, idx):
+        if not self.data_file:
+            self.read_data(self.path)
+        if isinstance(idx, int):
+            i = idx
+            self.check_index(i)
+            tensor_size = self.sizes[self.dim_offsets[i]:self.dim_offsets[i + 1]]
+            a = np.empty(tensor_size, dtype=self.dtype)
+            self.data_file.seek(self.data_offsets[i] * self.element_size)
+            self.data_file.readinto(a)
+            return a
+        elif isinstance(idx, slice):
+            start, stop, step = idx.indices(len(self))
+            if step != 1:
+                raise ValueError("Slices into indexed_dataset must be contiguous")
+            sizes = self.sizes[self.dim_offsets[start]:self.dim_offsets[stop]]
+            size = sum(sizes)
+            a = np.empty(size, dtype=self.dtype)
+            self.data_file.seek(self.data_offsets[start] * self.element_size)
+            self.data_file.readinto(a)
+            offsets = list(accumulate(sizes))
+            sents = np.split(a, offsets[:-1])
+            return sents
+
+    def __len__(self):
+        return self._len
+
+    def num_tokens(self, index):
+        return self.sizes[index]
+
+    def size(self, index):
+        return self.sizes[index]
+
+    @staticmethod
+    def exists(path):
+        return (
+            os.path.exists(index_file_path(path)) and os.path.exists(data_file_path(path))
+        )
+
+    @property
+    def supports_prefetch(self):
+        return False  # avoid prefetching to save memory
+
+
+class IndexedCachedDataset(IndexedDataset):
+
+    def __init__(self, path):
+        super().__init__(path)
+        self.cache = None
+        self.cache_index = {}
+
+    @property
+    def supports_prefetch(self):
+        return True
+
+    def prefetch(self, indices):
+        if all(i in self.cache_index for i in indices):
+            return
+        if not self.data_file:
+            self.read_data(self.path)
+        indices = sorted(set(indices))
+        total_size = 0
+        for i in indices:
+            total_size += self.data_offsets[i + 1] - self.data_offsets[i]
+        self.cache = np.empty(total_size, dtype=self.dtype)
+        ptx = 0
+        self.cache_index.clear()
+        for i in indices:
+            self.cache_index[i] = ptx
+            size = self.data_offsets[i + 1] - self.data_offsets[i]
+            a = self.cache[ptx: ptx + size]
+            self.data_file.seek(self.data_offsets[i] * self.element_size)
+            self.data_file.readinto(a)
+            ptx += size
+        if self.data_file:
+            # close and delete data file after prefetch so we can pickle
+            self.data_file.close()
+            self.data_file = None
+
+    # @lru_cache(maxsize=8)
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            i = idx
+            self.check_index(i)
+            tensor_size = self.sizes[self.dim_offsets[i]:self.dim_offsets[i + 1]]
+            a = np.empty(tensor_size, dtype=self.dtype)
+            ptx = self.cache_index[i]
+            np.copyto(a, self.cache[ptx: ptx + a.size])
+            return a
+        elif isinstance(idx, slice):
+            # Hack just to make this work, can optimizer later if necessary
+            sents = []
+            for i in range(*idx.indices(len(self))):
+                sents.append(self[i])
+            return sents
+
+
+class IndexedDatasetBuilder(object):
+    element_sizes = {
+        np.uint8: 1,
+        np.int8: 1,
+        np.int16: 2,
+        np.int32: 4,
+        np.int64: 8,
+        np.float: 4,
+        np.double: 8
+    }
+
+    def __init__(self, out_file, dtype=np.int32):
+        self.out_file = open(out_file, 'wb')
+        self.dtype = dtype
+        self.data_offsets = [0]
+        self.dim_offsets = [0]
+        self.sizes = []
+        self.element_size = self.element_sizes[self.dtype]
+        self.doc_idx = [0]
+
+    def add_item(self, tensor):
+        bytes = self.out_file.write(np.array(tensor.numpy(), dtype=self.dtype))
+        self.data_offsets.append(self.data_offsets[-1] + bytes / self.element_size)
+        for s in tensor.size():
+            self.sizes.append(s)
+        self.dim_offsets.append(self.dim_offsets[-1] + len(tensor.size()))
+
+    def end_document(self):
+        self.doc_idx.append(len(self.sizes))
+
+    def merge_file_(self, another_file):
+        index = IndexedDataset(another_file)
+        assert index.dtype == self.dtype
+
+        begin = self.data_offsets[-1]
+        for offset in index.data_offsets[1:]:
+            self.data_offsets.append(begin + offset)
+        self.sizes.extend(index.sizes)
+        begin = self.dim_offsets[-1]
+        for dim_offset in index.dim_offsets[1:]:
+            self.dim_offsets.append(begin + dim_offset)
+
+        with open(data_file_path(another_file), 'rb') as f:
+            while True:
+                data = f.read(1024)
+                if data:
+                    self.out_file.write(data)
+                else:
+                    break
+
+    def finalize(self, index_file):
+        self.out_file.close()
+        index = open(index_file, 'wb')
+        index.write(b'TNTIDX\x00\x00')
+        index.write(struct.pack('<Q', 1))
+        index.write(struct.pack('<QQ', code(self.dtype), self.element_size))
+        index.write(struct.pack('<QQ', len(self.data_offsets) - 1, len(self.sizes)))
+        index.write(struct.pack('<Q', len(self.doc_idx)))
+        write_longs(index, self.dim_offsets)
+        write_longs(index, self.data_offsets)
+        write_longs(index, self.sizes)
+        write_longs(index, self.doc_idx)
+        index.close()
+
+
+def _warmup_mmap_file(path):
+    with open(path, 'rb') as stream:
+        while stream.read(100 * 1024 * 1024):
+            pass
+
+
+class MMapIndexedDataset(torch.utils.data.Dataset):
+    class Index(object):
+        _HDR_MAGIC = b'MMIDIDX\x00\x00'
+
+        @classmethod
+        def writer(cls, path, dtype):
+            class _Writer(object):
+                def __enter__(self):
+                    self._file = open(path, 'wb')
+
+                    self._file.write(cls._HDR_MAGIC)
+                    self._file.write(struct.pack('<Q', 1))
+                    self._file.write(struct.pack('<B', code(dtype)))
+
+                    return self
+
+                @staticmethod
+                def _get_pointers(sizes):
+                    dtype_size = dtype().itemsize
+                    address = 0
+                    pointers = []
+
+                    for size in sizes:
+                        pointers.append(address)
+                        address += size * dtype_size
+
+                    return pointers
+
+                def write(self, sizes, doc_idx):
+                    pointers = self._get_pointers(sizes)
+
+                    self._file.write(struct.pack('<Q', len(sizes)))
+                    self._file.write(struct.pack('<Q', len(doc_idx)))
+
+                    sizes = np.array(sizes, dtype=np.int32)
+                    self._file.write(sizes.tobytes(order='C'))
+                    del sizes
+
+                    pointers = np.array(pointers, dtype=np.int64)
+                    self._file.write(pointers.tobytes(order='C'))
+                    del pointers
+
+                    doc_idx = np.array(doc_idx, dtype=np.int64)
+                    self._file.write(doc_idx.tobytes(order='C'))
+
+                def __exit__(self, exc_type, exc_val, exc_tb):
+                    self._file.close()
+
+            return _Writer()
+
+        def __init__(self, path, skip_warmup=False):
+            with open(path, 'rb') as stream:
+                magic_test = stream.read(9)
+                assert self._HDR_MAGIC == magic_test, (
+                    'Index file doesn\'t match expected format. '
+                    'Make sure that --dataset-impl is configured properly.'
+                )
+                version = struct.unpack('<Q', stream.read(8))
+                assert (1,) == version
+
+                dtype_code, = struct.unpack('<B', stream.read(1))
+                self._dtype = dtypes[dtype_code]
+                self._dtype_size = self._dtype().itemsize
+
+                self._len = struct.unpack('<Q', stream.read(8))[0]
+                self._doc_count = struct.unpack('<Q', stream.read(8))[0]
+                offset = stream.tell()
+
+            if not skip_warmup:
+                bmt.print_rank("    warming up index mmap file...")
+                _warmup_mmap_file(path)
+
+            self._bin_buffer_mmap = np.memmap(path, mode='r', order='C')
+            self._bin_buffer = memoryview(self._bin_buffer_mmap)
+            bmt.print_rank("    reading sizes...")
+            self._sizes = np.frombuffer(
+                self._bin_buffer,
+                dtype=np.int32,
+                count=self._len,
+                offset=offset)
+            bmt.print_rank("    reading pointers...")
+            self._pointers = np.frombuffer(self._bin_buffer, dtype=np.int64, count=self._len,
+                                           offset=offset + self._sizes.nbytes)
+            bmt.print_rank("    reading document index...")
+            self._doc_idx = np.frombuffer(self._bin_buffer, dtype=np.int64, count=self._doc_count,
+                                          offset=offset + self._sizes.nbytes + self._pointers.nbytes)
+
+        def __del__(self):
+            self._bin_buffer_mmap._mmap.close()
+            del self._bin_buffer_mmap
+
+        @property
+        def dtype(self):
+            return self._dtype
+
+        @property
+        def sizes(self):
+            return self._sizes
+
+        @property
+        def doc_idx(self):
+            return self._doc_idx
+
+        @lru_cache(maxsize=8)
+        def __getitem__(self, i):
+            return self._pointers[i], self._sizes[i]
+
+        def __len__(self):
+            return self._len
+
+    def __init__(self, path, skip_warmup=False):
+        super().__init__()
+
+        self._path = None
+        self._index = None
+        self._bin_buffer = None
+
+        self._do_init(path, skip_warmup)
+
+    def __getstate__(self):
+        return self._path
+
+    def __setstate__(self, state):
+        self._do_init(state)
+
+    def _do_init(self, path, skip_warmup):
+        self._path = path
+        self._index = self.Index(index_file_path(self._path), skip_warmup)
+
+        if not skip_warmup:
+            bmt.print_rank("    warming up data mmap file...")
+            _warmup_mmap_file(data_file_path(self._path))
+        bmt.print_rank("    creating numpy buffer of mmap...")
+        self._bin_buffer_mmap = np.memmap(data_file_path(self._path), mode='r', order='C')
+        bmt.print_rank("    creating memory view of numpy buffer...")
+        self._bin_buffer = memoryview(self._bin_buffer_mmap)
+
+    def __del__(self):
+        self._bin_buffer_mmap._mmap.close()
+        del self._bin_buffer_mmap
+        del self._index
+
+    def __len__(self):
+        return len(self._index)
+
+    # @lru_cache(maxsize=8)
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            assert idx < len(self._index), "Index {} out of range: {}".format(idx, len(self._index))
+            ptr, size = self._index[idx]
+            np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype,
+                                     count=size, offset=ptr)
+            return np_array
+        elif isinstance(idx, slice):
+            start, stop, step = idx.indices(len(self))
+            if step != 1:
+                raise ValueError("Slices into indexed_dataset must be contiguous")
+            ptr = self._index._pointers[start]
+            sizes = self._index._sizes[idx]
+            offsets = list(accumulate(sizes))
+            total_size = sum(sizes)
+            np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype,
+                                     count=total_size, offset=ptr)
+            sents = np.split(np_array, offsets[:-1])
+            return sents
+
+    def get(self, idx, offset=0, length=None):
+        """ Retrieves a single item from the dataset with the option to only
+        return a portion of the item.
+
+        get(idx) is the same as [idx] but get() does not support slicing.
+        """
+        ptr, size = self._index[idx]
+        if length is None:
+            length = size - offset
+        ptr += offset * np.dtype(self._index.dtype).itemsize
+        np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype,
+                                 count=length, offset=ptr)
+        return np_array
+
+    @property
+    def sizes(self):
+        return self._index.sizes
+
+    # @property
+    # def doc_idx(self):
+    #     return self._index.doc_idx
+
+    # def get_doc_idx(self):
+    #     return self._index._doc_idx
+
+    # def set_doc_idx(self, doc_idx_):
+    #     self._index._doc_idx = doc_idx_
+
+    @property
+    def supports_prefetch(self):
+        return False
+
+    @staticmethod
+    def exists(path):
+        return (
+            os.path.exists(index_file_path(path)) and os.path.exists(data_file_path(path))
+        )
+
+
+class MMapIndexedDatasetBuilder(object):
+    def __init__(self, out_file, dtype=np.int64):
+        self._data_file = open(out_file, 'wb')
+        self._dtype = dtype
+        self._sizes = []
+        self._doc_idx = [0]
+
+    def add_item(self, tensor):
+        np_array = np.array(tensor.numpy(), dtype=self._dtype)
+        self._data_file.write(np_array.tobytes(order='C'))
+        self._sizes.append(np_array.size)
+
+    def end_document(self):
+        self._doc_idx.append(len(self._sizes))
+
+    def merge_file_(self, another_file):
+        # Concatenate index
+        index = MMapIndexedDataset.Index(index_file_path(another_file))
+        assert index.dtype == self._dtype
+
+        for size in index.sizes:
+            self._sizes.append(size)
+
+        # Concatenate data
+        with open(data_file_path(another_file), 'rb') as f:
+            shutil.copyfileobj(f, self._data_file)
+
+    def finalize(self, index_file):
+        self._data_file.close()
+
+        with MMapIndexedDataset.Index.writer(index_file, self._dtype) as index:
+            index.write(self._sizes, self._doc_idx)
\ No newline at end of file
diff --git a/examples/t5/data/vocab_file b/examples/t5/data/vocab_file
new file mode 100644
index 0000000..4e28ff6
Binary files /dev/null and b/examples/t5/data/vocab_file differ
diff --git a/examples/t5/preprocess.py b/examples/t5/preprocess.py
new file mode 100644
index 0000000..409ec1e
--- /dev/null
+++ b/examples/t5/preprocess.py
@@ -0,0 +1,213 @@
+# coding=utf-8
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Processing data for pretraining."""
+
+import argparse
+import io
+import json
+import multiprocessing
+import os
+import sys
+import time
+import math
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+import random
+import numpy as np
+try:
+    import nltk
+    nltk_available = True
+except ImportError:
+    nltk_available = False
+
+from transformers import T5Tokenizer
+
+from data import indexed_dataset
+
+# from ray.util.multiprocessing.pool import Pool
+
+random.seed(233)
+np.random.seed(233)
+g = torch.manual_seed(233)
+torch.cuda.manual_seed_all(233)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+class IdentitySplitter(object):
+    def tokenize(self, *text):
+        return text
+
+class Encoder(object):
+    def __init__(self, args):
+        self.args = args
+
+    def initializer(self):
+        # Use Encoder class as a container for global data
+        Encoder.tokenizer = T5Tokenizer.from_pretrained("t5-base")
+        Encoder.splitter = IdentitySplitter()
+
+    def encode(self, line):
+        # end with <eod>
+        if len(line) > 5000000:
+            return None, None, None, 0
+        data = line.strip()
+        data = data.replace("<n>", "\n")
+        doc_ids = Encoder.tokenizer.encode(data)
+        if len(doc_ids) < 10:
+            return None, None, None, 0
+        # doc_ids.append(Encoder.tokenizer.eod_id)
+        span_start_ends = self.random_spans_noise_mask(
+            len(doc_ids), noisy_density=self.args.noisy_density, mean_noise_span_length=self.args.mean_noise_span_length)
+        start = 0
+        no_noise_tokens = []
+        noise_tokens = []
+        for i, span in enumerate(span_start_ends):
+            # "+[0]" placeholder for sentinel
+            no_noise_tokens.append(doc_ids[start:span[0]] + [0])
+            noise_tokens.append([0] + doc_ids[span[0]:span[1]]) # "[0]+" placeholder for sentinel
+            start = span[1]
+
+        target_offset_map = []
+        target_len = 0
+        for (k, context), target in zip(enumerate(no_noise_tokens), noise_tokens):
+            context[-1] = self.tokenizer.vocab_size + k
+            target[0] = self.tokenizer.vocab_size + k
+            target_offset_map.extend([target_len, len(target)])
+            target_len += len(target)
+
+        no_noise_tokens = [x for y in no_noise_tokens for x in y]
+
+        noise_tokens = [x for y in noise_tokens for x in y]
+
+        return no_noise_tokens, noise_tokens, target_offset_map, len(line)
+
+    def random_spans_noise_mask(self, length, noisy_density=0.15, mean_noise_span_length=10.0):
+        num_noise_tokens = round(length * noisy_density)
+        num_noise_tokens = min(max(num_noise_tokens, 1), length - 1)
+        num_noise_spans = round(num_noise_tokens / mean_noise_span_length)
+        num_noise_spans = max(num_noise_spans, 1)
+        num_nonnoise_tokens = length - num_noise_tokens
+
+        def random_segment(seq_length, num_segment):
+            x = (torch.arange(seq_length - 1) < (num_segment - 1)).long()
+            a = torch.randperm(seq_length - 1, generator=g)
+            x = x[a]
+            x = F.pad(x, [1, 0])
+            segment_id = torch.cumsum(x, dim=0)
+            segment_lengths = torch.zeros(num_segment, dtype=torch.long).scatter_add_(0, segment_id, torch.ones(seq_length, dtype=torch.long))
+
+            return segment_lengths
+
+        noise_span_lengths = random_segment(num_noise_tokens, num_noise_spans)
+        nonnoise_span_lengths = random_segment(num_nonnoise_tokens, num_noise_spans)
+        interleaved_span_lengths = torch.stack([nonnoise_span_lengths, noise_span_lengths], dim=1).view(num_noise_spans * 2)
+        span_start_ends = torch.cumsum(interleaved_span_lengths, dim=0).view(-1, 2)
+        return span_start_ends.tolist()
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    group = parser.add_argument_group(title='input data')
+    group.add_argument('--input', default="openwebtext.txt", type=str, help='Path to input TXT')
+    
+    group = parser.add_argument_group(title='output data')
+    group.add_argument("--output_path", default="pretrain_data", type=str)
+    group.add_argument('--output_prefix', default="webtext", type=str,
+                       help='Path to binary output file without suffix')
+    group.add_argument('--dataset_impl', type=str, default='mmap',
+                       choices=['lazy', 'cached', 'mmap'])
+
+    group = parser.add_argument_group(title='runtime')
+    group.add_argument('--workers', type=int, default=32,
+                       help='Number of worker processes to launch')
+    group.add_argument('--log_interval', type=int, default=10000,
+                       help='Interval between progress updates')
+    group.add_argument('--noisy_density', type=float, default=0.15)
+    group.add_argument('--mean_noise_span_length', type=int, default=3)
+
+    args = parser.parse_args()
+    args.keep_empty = False
+
+    args.rank = 0
+    args.make_vocab_size_divisible_by = 128
+
+    return args
+
+def main():
+    args = get_args()
+    startup_start = time.time()
+
+    print("Opening", args.input)
+
+    fin = open(args.input, 'r')
+    encoder = Encoder(args)
+    tokenizer = T5Tokenizer.from_pretrained("t5-base")
+    pool = multiprocessing.Pool(args.workers, initializer=encoder.initializer)
+    
+    # use the tokenizer to encode the sentences
+    encoded_docs = pool.imap_unordered(encoder.encode, fin, 30)
+
+    level = "document"
+
+    print(f"Vocab size: {tokenizer.vocab_size}")
+    print(f"Output prefix: {args.output_prefix}")
+    context_bin_file = os.path.join(args.output_path, "{}_{}_context.bin".format(args.output_prefix, level))
+    context_idx_file = os.path.join(args.output_path,  "{}_{}_context.idx".format(args.output_prefix, level))
+    target_bin_file = os.path.join(args.output_path,  "{}_{}_target.bin".format(args.output_prefix, level))
+    target_idx_file = os.path.join(args.output_path,  "{}_{}_target.idx".format(args.output_prefix, level))
+    target_offset_bin_file = os.path.join(args.output_path,  "{}_{}_target_offset.bin".format(args.output_prefix, level))
+    target_offset_idx_file = os.path.join(args.output_path,  "{}_{}_target_offset.idx".format(args.output_prefix, level))
+    
+
+    # set dtype == int64
+    # max idx: 2 ** 63 - 1
+    builder_context = indexed_dataset.make_builder(context_bin_file, impl=args.dataset_impl, dtype=np.int32)
+    builder_target = indexed_dataset.make_builder(target_bin_file, impl=args.dataset_impl, dtype=np.int32)
+    builder_target_offset = indexed_dataset.make_builder(target_offset_bin_file, impl=args.dataset_impl, dtype=np.int32)
+
+    startup_end = time.time()
+    proc_start = time.time()
+    total_bytes_processed = 0
+    print("Time to startup:", startup_end - startup_start)
+
+    # sentinel_idx = tokenizer.vocab_size # start from the last token of the tokenizer
+    print("tokenizer vocab size:", tokenizer.vocab_size)
+    for i, (no_noise_tokens, noise_tokens, target_offset_map, bytes_processed) in enumerate(encoded_docs, start=1):
+        if no_noise_tokens is None or noise_tokens is None:
+            continue
+        total_bytes_processed += bytes_processed
+
+        builder_context.add_item(torch.IntTensor(no_noise_tokens))
+        builder_target.add_item(torch.IntTensor(noise_tokens))
+        builder_target_offset.add_item(torch.IntTensor(target_offset_map))
+        
+        if i % args.log_interval == 0:
+            current = time.time()
+            elapsed = current - proc_start
+            mbs = total_bytes_processed / elapsed / 1024 / 1024
+            print(f"Processed {i} documents",
+                  f"({i/elapsed} docs/s, {mbs} MB/s).",
+                  file=sys.stderr)
+
+    builder_context.finalize(context_idx_file)
+    builder_target.finalize(target_idx_file)
+    builder_target_offset.finalize(target_offset_idx_file)
+
+    pool.close()
+
+if __name__ == '__main__':
+    main()
diff --git a/setup.py b/setup.py
new file mode 100644
index 0000000..b86ae61
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,20 @@
+from setuptools import setup, find_packages
+
+def main():
+    setup(
+        name='bmcook',
+        version='0.1.1.1',
+        description="Model Compression for Big Models",
+        author="Baitao Gong",
+        author_email="gongbaitao11@gmail.com",
+        packages=find_packages(exclude='cpm_live'),
+        url="https://github.com/OpenBMB/BMCook",
+        install_requires=[
+            "bmtrain",
+            "model_center"
+        ],
+        license="Apache 2.0"
+    )
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file


Test patch:
diff --git a/examples/cpm_live/cpm_ant_test.py b/examples/cpm_live/cpm_ant_test.py
new file mode 100644
index 0000000..96a82d7
--- /dev/null
+++ b/examples/cpm_live/cpm_ant_test.py
@@ -0,0 +1,83 @@
+import torch
+import bmtrain as bmt
+from bmcook.utils.config import ConfigParser
+from bmcook import CPMAntTrainer
+
+import sys, os
+sys.path.insert(0, os.getcwd())
+from cpm_live.arguments import parse_args
+from cpm_live.models import CPMAntConfig, CPMAnt
+from cpm_live.tokenizers import CPMAntTokenizer
+from cpm_live.data import CPMAnt_Dataset, DistributedMMapIndexedDataset, BatchPacker
+
+def main():
+
+    bmt.init_distributed(seed = 233, loss_scale_factor = 2, loss_scale_steps = 512)
+
+    args = parse_args()
+    batch_size = 4
+
+    config = CPMAntConfig.from_json_file(args.model_config)
+    model = CPMAnt(config)
+    bmt.load(model, args.load)
+
+    teacher_config = CPMAntConfig.from_json_file(args.teacher_config)
+    teacher = CPMAnt(teacher_config)
+    bmt.load(teacher, args.load_teacher)
+
+    config = ConfigParser(args.cook_config)
+    
+
+    tokenizer = CPMAntTokenizer('vocabs/ant.txt')
+    dataset = CPMAnt_Dataset(
+        DistributedMMapIndexedDataset(args.data_path, "cpm_live_text_merge_context", bmt.rank(), bmt.world_size()),
+        max_length = 512 - 32, 
+        prompt_length = 32,
+        tokenizer = tokenizer
+    )
+    dataloader = BatchPacker(dataset, 512, batch_size=batch_size)
+    optimizer = bmt.optim.AdamOffloadOptimizer(model.parameters(), 
+                                                weight_decay=0.01, 
+                                                scale=1048576)
+    lr_scheduler = bmt.lr_scheduler.Noam(optimizer, 
+                                            start_lr = 0.1,
+                                            warmup_iter = 2000, 
+                                            end_iter = None,
+                                            num_iter = 0)
+    loss_func = bmt.loss.FusedCrossEntropy(ignore_index=-100)
+
+    CPMAntTrainer.set_compression(config, model, optimizer, teacher)
+
+    for iteration, data in enumerate(dataloader):
+
+        assert len(data["ctx"]) == batch_size
+        input_idx = data["ctx"].int().cuda()
+        input_length = data["len_ctx"].int().cuda()
+        input_context = data["context"].bool().cuda()
+        input_position = data["position"].float().cuda()
+        input_segment = data["segment"].int().cuda()
+        input_span = data["span"].int().cuda()
+        targets = data["tgt"].long().cuda()
+        task_info = data["task"].long().cuda()
+        
+        # ===========
+        optimizer.zero_grad()
+        torch.cuda.empty_cache()
+
+        # ===========
+        outputs = CPMAntTrainer.forward(model, loss_func, targets, input_idx, input_length, input_context, input_position, input_segment, input_span)
+        loss, lag_loss, sparsity, d_loss = outputs.loss, outputs.lag_loss, outputs.sparsity, outputs.d_loss
+        global_loss = bmt.sum_loss(loss - d_loss).item()
+
+        # ===========
+        loss = optimizer.loss_scale(loss) + lag_loss
+        loss.backward()
+        
+        # ===========
+        grad_norm = bmt.optim.clip_grad_norm(optimizer.param_groups, 4.0, scale = optimizer.scale, norm_type = 2)    
+        bmt.optim_step(optimizer, lr_scheduler)
+
+        bmt.print_rank('iter: {} | loss: {:.4f} | grad_norm: {} | d_loss: {} | lag_loss: {:.4f} | sparsity: {:.4f}'.format(iteration, global_loss, grad_norm, d_loss, lag_loss, sparsity))
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/examples/cpm_live/cpm_ant_test.sh b/examples/cpm_live/cpm_ant_test.sh
new file mode 100644
index 0000000..dd5ec97
--- /dev/null
+++ b/examples/cpm_live/cpm_ant_test.sh
@@ -0,0 +1,8 @@
+export CUDA_VISIBLE_DEVICES=2,3
+torchrun --nnodes=1 --nproc_per_node=2 --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=localhost:23333 cpm_live/cpm_ant_test.py \
+    --model-config ... \
+    --load ... \
+    --teacher-config ... \
+    --load-teacher ... \
+    --data-path ... \
+    --cook-config cpm_live/configs/prune.json \
\ No newline at end of file
diff --git a/gpt-example/train.py b/examples/gpt/gpt2_test.py
similarity index 56%
rename from gpt-example/train.py
rename to examples/gpt/gpt2_test.py
index fd6975c..98b6794 100644
--- a/gpt-example/train.py
+++ b/examples/gpt/gpt2_test.py
@@ -1,22 +1,14 @@
 import os
-import sys
-sys.path.insert(0, os.path.abspath('../'))
-
-
 import json
-import types
-from bmcook.moe import BMMoE
-from bmcook.quant import BMQuant
 import torch
 import random
-import bmtrain as bmt
 import time
+import bmtrain as bmt
 from data import MMapIndexedDataset, Dataset
-from bmcook.pruning import BMPrune
-from bmcook.distilling import BMDistill
 from bmcook.utils.arguments import parse_args
+from bmcook.utils.config import ConfigParser
+from bmcook import CookTrainer
 from pathlib import Path
-import json
 
 def print_inspect(model, name):
     bmt.print_rank(
@@ -25,30 +17,7 @@ def print_inspect(model, name):
         )
     )
 
-class Trainer:
-    @staticmethod
-    def batch_iter_shuf(dataset : Dataset, batch_size, rank, world_size):
-        st = 0
-        end = len(dataset)
-
-        local_size = len(dataset) // world_size
-        idx = list(range(local_size))
-        random.shuffle(idx)
-
-        batch = []
-        while st < local_size:
-            it = dataset[idx[st]*world_size + rank]
-            if it is not None:
-                batch.append( it )
-            st += 1
-            if len(batch) == batch_size:
-                yield {
-                    "ctx": torch.stack([it["ctx"] for it in batch]),
-                    "len_ctx": torch.LongTensor([it["len_ctx"] for it in batch]),
-                    'target': torch.stack([it["target"] for it in batch]),
-                }
-                batch = []
-
+class Dataloader:
     @staticmethod
     def batch_iter(dataset : Dataset, batch_size, rank, world_size):
         st = 0
@@ -67,25 +36,14 @@ def batch_iter(dataset : Dataset, batch_size, rank, world_size):
                 }
                 batch = []
 
-    @staticmethod
-    def forward(model, enc_input, enc_length, dec_input, dec_length, targets, loss_func):
-        outputs = model(
-            enc_input, enc_length, dec_input, dec_length, return_logits=True)
-        logits = outputs
-        batch, seq_len, vocab_out_size = logits.size()
-
-        loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))
-
-        return [loss, logits]
-
 from model_center.model import GPT2Config, GPT2
 
 config_map = {
-    'gpt2-base': GPT2Config
+    'gpt2-large': GPT2Config
 }
 
 model_map = {
-    'gpt2-base': GPT2
+    'gpt2-large': GPT2
 }
 
 def main():
@@ -102,14 +60,6 @@ def main():
     # teacher model has the same config as the student model
     teacher = model_map[args.model].from_pretrained(args.model, config=model_config)
 
-    def new_forward(model_self, enc_input, enc_length, dec_input, dec_length, return_logits=False):
-        return model_self.forward_old(dec_input, dec_length, return_logits=return_logits)
-    
-    model.forward_old = model.forward
-    model.forward = types.MethodType(new_forward, model)
-    teacher.forward_old = teacher.forward
-    teacher.forward = types.MethodType(new_forward, teacher)
-
     bmt.synchronize()
 
     # data
@@ -120,45 +70,8 @@ def new_forward(model_self, enc_input, enc_length, dec_input, dec_length, return
     optimizer = bmt.optim.AdamOptimizer(model.parameters(), scale=2**20)
     lr_scheduler = bmt.lr_scheduler.Noam(optimizer, start_lr=args.start_lr, warmup_iter=2000, end_iter=100000)
 
-    # bmcook config
-    from bmcook.utils.config import ConfigParser
     config = ConfigParser(args.cook_config)
-
-    # remove checkpointing
-    for _, v in model.named_modules():
-
-        if isinstance(v, bmt.TransformerBlockList):
-
-            def new_func(list_self, hidden_states, *args):
-                for i in range(len(list_self._modules)):
-                    hidden_states = list_self._modules[str(i)](hidden_states, *args)
-                return hidden_states
-
-            v.forward = types.MethodType(new_func, v)
-
-            for k in v._modules.keys():
-                state_dict = v._modules[k].state_dict()
-                for kk, vv in v._modules[k]._module.named_modules():
-                    if kk+'.weight' in state_dict:
-                        vv.weight.data = state_dict[kk+'.weight'].clone().cuda()
-                    if kk+'.bias' in state_dict:
-                        vv.bias.data = state_dict[kk+'.bias'].clone().cuda()
-                v._modules[k] = v._modules[k]._module
-    
-    # for distillation
-    Trainer.forward = BMDistill.set_forward(model, teacher, Trainer.forward, config)
-
-    # for pruning
-    BMPrune.compute_mask(model, config)
-    BMPrune.set_optim_for_pruning(optimizer)
-
-    # for quantization
-    BMQuant.quantize(model, config)
-
-    # for moefication
-    Trainer.forward = BMMoE.get_hidden(model, config, Trainer.forward)
-
-    bmt.synchronize()
+    CookTrainer.set_compression(config, model, optimizer, teacher)
     
     average_time = 0
     average_time_shift = 0.9
@@ -172,7 +85,7 @@ def new_func(list_self, hidden_states, *args):
         os.makedirs(save_dir / 'hiddens', exist_ok=True)
         model.eval()
 
-        for iteration, data in enumerate(Trainer.batch_iter(dataset, batch_size, bmt.rank(), bmt.world_size())):
+        for iteration, data in enumerate(Dataloader.batch_iter(dataset, batch_size, bmt.rank(), bmt.world_size())):
 
             if iteration == 100:
                 break
@@ -187,7 +100,7 @@ def new_func(list_self, hidden_states, *args):
             dec_length = dec_length.cuda()
             
             with torch.no_grad():
-                outputs = Trainer.forward(model, None, None, dec_input, dec_length, targets, loss_func)
+                outputs = CookTrainer.forward(model, loss_func, targets, dec_input, dec_length)
             
             torch.save(outputs[-1], save_dir / 'hiddens' / '{}_{}'.format(iteration, bmt.rank()))
                
@@ -204,7 +117,7 @@ def new_func(list_self, hidden_states, *args):
 
     for epoch in range(3):
         
-        for iteration, data in enumerate(Trainer.batch_iter(dataset, batch_size, bmt.rank(), bmt.world_size())):
+        for iteration, data in enumerate(Dataloader.batch_iter(dataset, batch_size, bmt.rank(), bmt.world_size())):
 
             st = time.time()
             optimizer.zero_grad()
@@ -218,14 +131,15 @@ def new_func(list_self, hidden_states, *args):
             dec_input = dec_input.cuda()
             dec_length = dec_length.cuda()
 
-            outputs = Trainer.forward(model, None, None, dec_input, dec_length, targets, loss_func)
+            outputs = CookTrainer.forward(model, loss_func, targets, dec_input, dec_length)
 
-            loss = outputs[0]
+            loss = outputs.loss
+            lag_loss, sparsity = bmt.sum_loss(outputs.lag_loss).item(), outputs.sparsity
             global_loss = bmt.sum_loss(loss).item()
-            loss = optimizer.loss_scale(loss)
+            loss = optimizer.loss_scale(loss) + outputs.lag_loss
 
             if do_distill:
-                distill_loss = bmt.sum_loss(outputs[-1]).item()
+                distill_loss = bmt.sum_loss(outputs.d_loss).item()
             else:
                 distill_loss = 0
 
@@ -240,13 +154,15 @@ def new_func(list_self, hidden_states, *args):
                 iteration_time = time.time() - st
                 average_time = average_time * average_time_shift + (1 - average_time_shift) * iteration_time
                 bmt.print_rank(
-                    "| Iter: {:6d} | loss: {:.4f} | kd_loss: {:.4f} | lr: {:.4e}, scale: {:10.4f} | time: {:.4f}".format(
+                    "| Iter: {:6d} | loss: {:.4f} | kd_loss: {:.4f} | lr: {:.4e}, scale: {:10.4f} | time: {:.4f}  | lag_loss: {:.4f} | sparsity: {:.4f}".format(
                         iteration,
                         global_loss-distill_loss,
                         distill_loss,
                         lr_scheduler.current_lr,
                         int(optimizer.scale),
-                        average_time / (1 - pow(average_time_shift, iteration + 1))
+                        average_time / (1 - pow(average_time_shift, iteration + 1)),
+                        lag_loss,
+                        sparsity
                     )
                 )
             
diff --git a/examples/gpt/gpt2_test.sh b/examples/gpt/gpt2_test.sh
new file mode 100644
index 0000000..42fc46a
--- /dev/null
+++ b/examples/gpt/gpt2_test.sh
@@ -0,0 +1,6 @@
+export CUDA_VISIBLE_DEVICES=2,3
+torchrun --nnodes=1 --nproc_per_node=2 --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=localhost gpt/gpt2_test.py \
+    --model gpt2-large \
+    --save-dir ... \
+    --data-path ... \
+    --cook-config gpt/configs/gpt2-prune.json \
\ No newline at end of file
diff --git a/examples/t5/t5_test.py b/examples/t5/t5_test.py
new file mode 100644
index 0000000..13329a3
--- /dev/null
+++ b/examples/t5/t5_test.py
@@ -0,0 +1,203 @@
+import os
+import json
+import torch
+import random
+import time
+import bmtrain as bmt
+from data import MMapIndexedDataset, Dataset
+from bmcook import CookTrainer
+from bmcook.utils.config import ConfigParser
+from bmcook.utils.arguments import parse_args
+from pathlib import Path
+
+def print_inspect(model, name):
+    bmt.print_rank(
+        bmt.inspect.format_summary(
+            bmt.inspect.inspect_model(model, name)
+        )
+    )
+
+class Dataloader:
+    @staticmethod
+    def batch_iter(dataset : Dataset, batch_size, rank, world_size):
+        st = 0
+        end = len(dataset)
+        batch = []
+        while st < end:
+            it = dataset[st + rank]
+            if it is not None:
+                batch.append( it )
+            st += world_size
+            if len(batch) == batch_size:
+                yield {
+                    "ctx_context": torch.stack([it["ctx_context"] for it in batch]),
+                    "len_ctx_context": torch.LongTensor([it["len_ctx_context"] for it in batch]),
+                    "ctx_target": torch.stack([it["ctx_target"] for it in batch]),
+                    "len_ctx_target": torch.LongTensor([it["len_ctx_target"] for it in batch]),
+                    'target': torch.stack([it["target"] for it in batch]),
+                }
+                batch = []
+
+from model_center.model import GPT2Config, GPT2, T5Config, T5
+
+config_map = {
+    'gpt2-base': GPT2Config,
+    't5-3b': T5Config,
+    't5-large': T5Config,
+    't5-base': T5Config,
+}
+
+model_map = {
+    'gpt2-base': GPT2,
+    't5-3b': T5,
+    't5-large': T5,
+    't5-base': T5,
+}
+
+def main():
+    bmt.init_distributed()
+
+    args = parse_args()
+    save_dir = Path(args.save_dir)
+    ckpt_dir = save_dir / 'checkpoints'
+    os.makedirs(ckpt_dir, exist_ok=True)
+    json.dump(vars(args), open(save_dir / 'train_args.json', 'w'), indent=2)
+
+    model_config = config_map[args.model].from_pretrained(args.model)
+    model = model_map[args.model].from_pretrained(args.model, config=model_config)
+    
+    # teacher model has the same config as the student model
+    teacher = model_map[args.model].from_pretrained(args.model, config=model_config)
+
+    bmt.synchronize()
+
+    # data
+    batch_size = 4
+    ctx_len = 512
+    tar_len = 512
+
+    loss_func = torch.nn.CrossEntropyLoss(ignore_index=-100)
+    optimizer = bmt.optim.AdamOptimizer(model.parameters(), scale=2**20)
+    lr_scheduler = bmt.lr_scheduler.Noam(optimizer, start_lr=args.start_lr, warmup_iter=2000, end_iter=100000)
+
+    config = ConfigParser(args.cook_config)
+    CookTrainer.set_compression(config, model, optimizer, teacher)    
+    
+    average_time = 0
+    average_time_shift = 0.9
+
+    dataset = Dataset(
+        MMapIndexedDataset(os.path.join(args.data_path, 'webtext_document_context')),
+        MMapIndexedDataset(os.path.join(args.data_path, 'webtext_document_target')),
+        ctx_len,
+        tar_len
+    )
+
+
+    if config.get('MoEfication')['is_moefy']:
+        os.makedirs(save_dir / 'hiddens', exist_ok=True)
+        model.eval()
+
+        for iteration, data in enumerate(Dataloader.batch_iter(dataset, batch_size, bmt.rank(), bmt.world_size())):
+
+            if iteration == 100:
+                break
+
+            enc_input = data["ctx_context"].int()
+            enc_length = data["len_ctx_context"].int()
+
+            enc_input = enc_input.cuda()
+            enc_length = enc_length.cuda()
+
+            dec_input = data["ctx_target"].int()
+            dec_length = data["len_ctx_target"].int()
+            dec_mask = torch.arange(tar_len)[None, :].repeat(batch_size, 1) < dec_length[:, None]
+            targets = torch.where(dec_mask, data["target"].long(), torch.scalar_tensor(-100, dtype=torch.long))
+
+            targets = targets.cuda()
+            dec_input = dec_input.cuda()
+            dec_length = dec_length.cuda()
+            
+            with torch.no_grad():
+                outputs = CookTrainer.forward(model, loss_func, targets, enc_input, enc_length, dec_input, dec_length)
+            
+            torch.save(outputs[-1], 'results/hiddens/' + '{}_{}'.format(iteration, bmt.rank()))
+               
+            bmt.print_rank("Iteration:", iteration)
+        exit()
+
+    do_distill = True
+    distill_config = config.get('distillation')
+    if distill_config['ce_scale'] + distill_config['mse_hidn_scale'] + distill_config['mse_att_scale'] == 0:
+        do_distill = False
+
+    model.train()
+    teacher.eval()
+
+    for epoch in range(3):
+        
+        for iteration, data in enumerate(Dataloader.batch_iter(dataset, batch_size, bmt.rank(), bmt.world_size())):
+
+            st = time.time()
+            optimizer.zero_grad()
+
+            enc_input = data["ctx_context"].int()
+            enc_length = data["len_ctx_context"].int()
+
+            enc_input = enc_input.cuda()
+            enc_length = enc_length.cuda()
+
+            dec_input = data["ctx_target"].int()
+            dec_length = data["len_ctx_target"].int()
+            dec_mask = torch.arange(tar_len)[None, :].repeat(batch_size, 1) < dec_length[:, None]
+            targets = torch.where(dec_mask, data["target"].long(), torch.scalar_tensor(-100, dtype=torch.long))
+
+            targets = targets.cuda()
+            dec_input = dec_input.cuda()
+            dec_length = dec_length.cuda()
+
+            outputs = CookTrainer.forward(model, loss_func, targets, enc_input, enc_length, dec_input, dec_length)
+
+            loss = outputs.loss
+            lag_loss, sparsity = bmt.sum_loss(outputs.lag_loss).item(), outputs.sparsity
+            
+            global_loss = bmt.sum_loss(loss).item()
+            loss = optimizer.loss_scale(loss) + outputs.lag_loss
+
+            if do_distill:
+                distill_loss = bmt.sum_loss(outputs.d_loss).item()
+            else:
+                distill_loss = 0
+            
+            loss.backward()
+
+            bmt.optim_step(optimizer, lr_scheduler)
+            
+
+            if iteration % args.log_interval == 0:
+                iteration_time = time.time() - st
+                average_time = average_time * average_time_shift + (1 - average_time_shift) * iteration_time
+                bmt.print_rank(
+                    "| Iter: {:6d} | loss: {:.4f} | kd_loss: {:.4f} | lr: {:.4e}, scale: {:10.4f} | time: {:.4f} | lag_loss: {:.4f} | sparsity: {:.4f}".format(
+                        iteration,
+                        global_loss-distill_loss,
+                        distill_loss,
+                        lr_scheduler.current_lr,
+                        int(optimizer.scale),
+                        average_time / (1 - pow(average_time_shift, iteration + 1)),
+                        lag_loss,
+                        sparsity,
+                    )
+                )
+            
+            if iteration % args.save_interval == 0: 
+                ckpt_file = Path(args.save_dir, 'checkpoints', f'ckpt-{iteration}.pt')
+                bmt.save(model, ckpt_file) 
+
+    ckpt_file = Path(args.save_dir, 'checkpoint.pt') 
+    bmt.save(model, ckpt_file) 
+    bmt.synchronize()
+
+
+if __name__ == '__main__': 
+    main()
diff --git a/examples/t5/t5_test.sh b/examples/t5/t5_test.sh
new file mode 100644
index 0000000..60ff5a0
--- /dev/null
+++ b/examples/t5/t5_test.sh
@@ -0,0 +1,6 @@
+export CUDA_VISIBLE_DEVICES=2,3
+torchrun --nnodes=1 --nproc_per_node=2 --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=localhost t5/t5_test.py \
+    --model t5-large \
+    --save-dir ... \
+    --data-path ... \
+    --cook-config t5/configs/t5-combine.json \
\ No newline at end of file