Merge branch 'develop' into feat/refactor-impl

Author: zigzagcai

configs/7B_isp_sft.py

@@ -243,6 +243,7 @@
 cudnn_deterministic = False
 cudnn_benchmark = False
 
+
 monitor = dict(
     # feishu alert configs
     alert=dict(

internlm/model/model_implementations/transformers/modeling_internlm.py

@@ -18,7 +18,6 @@
 from internlm.core.context import ParallelMode
 from internlm.core.context import global_context as gpc
 from internlm.core.naive_amp import set_output_attr_to_module
-from internlm.core.parallel.shard import partition_uniform
 from internlm.model.model_implementations.transformers.base_model import (
     BaseTransformerModel,
 )

internlm/model/model_ops/moe/moe.py

@@ -1,5 +1,4 @@
 import torch
-import torch.nn.functional as F
 
 from internlm.core.context import ParallelMode
 from internlm.core.context import global_context as gpc

internlm/model/model_ops/ops/ring_flash_attn/zigzag_ring_flash_attn_with_sliding_window.py

@@ -5,11 +5,10 @@
 from flash_attn.flash_attn_interface import _flash_attn_backward, _flash_attn_forward
 
 from internlm.core.context import global_context as gpc
+from internlm.core.parallel.comm import get_offload_manager
 
 from .utils import RingComm, update_out_and_lse
 
-fa_output_mapping = {}
-
 
 def create_buffer(tensor):
     buffer_shape = list(tensor.shape)
@@ -443,9 +442,10 @@ def forward(
         k = k.contiguous()
         v = v.contiguous()
 
-        if gpc.is_forward is False and gpc.config.selective_checkpoint:
-            assert layer_idx in fa_output_mapping
-            out, softmax_lse = fa_output_mapping.pop(layer_idx)
+        _ckpt_block_num = int(gpc.config.model.checkpoint * gpc.config.isp_num_layers)
+
+        if gpc.is_forward is False and gpc.config.selective_checkpoint and layer_idx < _ckpt_block_num:
+            out, softmax_lse = get_offload_manager().get_fa_output_with_layer(layer_idx)
         else:
             out, softmax_lse = zigzag_double_ring_flash_attn_forward(
                 context_group,
@@ -460,8 +460,8 @@ def forward(
             )
 
         # store attn forward output to avoid re-computation of attn when activation checkpoint is enabled
-        if gpc.is_forward and gpc.config.selective_checkpoint:
-            fa_output_mapping[layer_idx] = (out, softmax_lse)
+        if gpc.is_forward and gpc.config.selective_checkpoint and layer_idx < _ckpt_block_num:
+            get_offload_manager().insert_fa_output_with_layer(layer_idx=layer_idx, output=(out, softmax_lse))
 
         # this should be out_padded
         ctx.save_for_backward(q, k, v, out, softmax_lse)

internlm/model/model_ops/utils.py

@@ -5,7 +5,6 @@
 from tqdm import tqdm
 
 from internlm.core.context import global_context as gpc
-from internlm.model.model_ops.modules.mha import MHA
 from internlm.utils.logger import get_logger
 from internlm.utils.storage_manager import get_fns, llm_load
 from internlm.utils.utils import TensorParallelMode