Add fine-tuning code

modeling_InternLM.py

@@ -6,8 +6,6 @@ import torch
 import torch.utils.checkpoint
 import torch.utils.checkpoint
 from einops import rearrange
-#import rotary_emb
-#from flash_attn.layers.rotary import ApplyRotaryEmbQKV_ as LegacyApplyRotaryEmbQKV_
 from torch import nn
 from torch.nn import CrossEntropyLoss
 from transformers.activations import ACT2FN
@@ -18,14 +16,17 @@ from transformers.utils import logging
 from .configuration_InternLM_XComposer import InternLMXComposerConfig
 from .modeling_utils import LoRALinear
 
+try:
+    import rotary_emb
+except Exception as e:
+    print('Please following docs/install.md to install rotary_emb if you want to do fine-tuning')
+
 logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "InternLMXComposerConfig"
 
 
 def rotary_embed(x1, x2, cos, sin, conj):
-    # print(x1.shape, x2.shape, cos.shape, sin.shape)
-    #[5, 1, 32, 64] [1, 1, 64]
     x1, x2 = x1.float(), x2.float()
     if conj:
         x1, x2 = x1 * cos + x2 * sin, x1 * sin + x2 * cos
@@ -35,7 +36,6 @@ def rotary_embed(x1, x2, cos, sin, conj):
 
 
 class LegacyApplyRotaryEmbQKV_(torch.autograd.Function):
-
     @staticmethod
     def forward(ctx, qkv, cos, sin, cos_k=None, sin_k=None, interleaved=False):
         """
@@ -55,18 +55,26 @@ class LegacyApplyRotaryEmbQKV_(torch.autograd.Function):
         assert seqlen <= rotary_seqlen
         cos_k = cos if cos_k is None else cos_k
         sin_k = sin if sin_k is None else sin_k
-        assert sin.shape == cos_k.shape == sin_k.shape == (rotary_seqlen, rotary_dim // 2)
+        assert sin.shape == cos_k.shape == sin_k.shape == (rotary_seqlen,
+                                                           rotary_dim // 2)
         q_ro = qkv[:, :, 0, :, :rotary_dim]
-        q1, q2 = q_ro.chunk(2, dim=-1) if not interleaved else (q_ro[..., ::2], q_ro[..., 1::2])
+        q1, q2 = q_ro.chunk(2, dim=-1) if not interleaved else (q_ro[..., ::2],
+                                                                q_ro[...,
+                                                                     1::2])
         # rotary_emb.apply_rotary(q1, q2, rearrange(cos[:seqlen], 's d -> s 1 d'),
         #                         rearrange(sin[:seqlen], 's d -> s 1 d'), q1, q2, False)
-        q1, q2 = rotary_embed(q1, q2, rearrange(cos[:seqlen], 's d -> s 1 d'), rearrange(sin[:seqlen], 's d -> s 1 d'), False)
+        q1, q2 = rotary_embed(q1, q2, rearrange(cos[:seqlen], 's d -> s 1 d'),
+                              rearrange(sin[:seqlen], 's d -> s 1 d'), False)
         qkv[:, :, 0, :, :rotary_dim] = torch.cat([q1, q2], dim=-1)
         k_ro = qkv[:, :, 1, :, :rotary_dim]
-        k1, k2 = k_ro.chunk(2, dim=-1) if not interleaved else (k_ro[..., ::2], k_ro[..., 1::2])
+        k1, k2 = k_ro.chunk(2, dim=-1) if not interleaved else (k_ro[..., ::2],
+                                                                k_ro[...,
+                                                                     1::2])
         # rotary_emb.apply_rotary(k1, k2, rearrange(cos_k[:seqlen], 's d -> s 1 d'),
         #                         rearrange(sin_k[:seqlen], 's d -> s 1 d'), k1, k2, False)
-        k1, k2 = rotary_embed(k1, k2, rearrange(cos_k[:seqlen], 's d -> s 1 d'), rearrange(sin_k[:seqlen], 's d -> s 1 d'), False)
+        k1, k2 = rotary_embed(k1, k2, rearrange(cos_k[:seqlen],
+                                                's d -> s 1 d'),
+                              rearrange(sin_k[:seqlen], 's d -> s 1 d'), False)
         qkv[:, :, 1, :, :rotary_dim] = torch.cat([k1, k2], dim=-1)
         ctx.save_for_backward(cos, sin, cos_k, sin_k)
         ctx.interleaved = interleaved
@@ -79,18 +87,69 @@ class LegacyApplyRotaryEmbQKV_(torch.autograd.Function):
         rotary_dim = cos.shape[-1]
         rotary_dim *= 2
         dq_ro = dqkv[:, :, 0, :, :rotary_dim]
-        dq1, dq2 = (dq_ro.chunk(2, dim=-1) if not ctx.interleaved
-                    else (dq_ro[..., ::2], dq_ro[..., 1::2]))
-        rotary_emb.apply_rotary(dq1, dq2, rearrange(cos[:seqlen], 's d -> s 1 d'),
-                                rearrange(sin[:seqlen], 's d -> s 1 d'), dq1, dq2, True)
+        dq1, dq2 = (dq_ro.chunk(2, dim=-1) if not ctx.interleaved else
+                    (dq_ro[..., ::2], dq_ro[..., 1::2]))
+        rotary_emb.apply_rotary(dq1, dq2,
+                                rearrange(cos[:seqlen], 's d -> s 1 d'),
+                                rearrange(sin[:seqlen], 's d -> s 1 d'), dq1,
+                                dq2, True)
         dk_ro = dqkv[:, :, 1, :, :rotary_dim]
-        dk1, dk2 = (dk_ro.chunk(2, dim=-1) if not ctx.interleaved
-                    else (dk_ro[..., ::2], dk_ro[..., 1::2]))
-        rotary_emb.apply_rotary(dk1, dk2, rearrange(cos_k[:seqlen], 's d -> s 1 d'),
-                                rearrange(sin_k[:seqlen], 's d -> s 1 d'), dk1, dk2, True)
+        dk1, dk2 = (dk_ro.chunk(2, dim=-1) if not ctx.interleaved else
+                    (dk_ro[..., ::2], dk_ro[..., 1::2]))
+        rotary_emb.apply_rotary(dk1, dk2,
+                                rearrange(cos_k[:seqlen], 's d -> s 1 d'),
+                                rearrange(sin_k[:seqlen], 's d -> s 1 d'), dk1,
+                                dk2, True)
         return dqkv, None, None, None, None, None
 
 
+class ApplyRotaryEmbQKV_(torch.autograd.Function):
+    """
+    ApplyRotaryEmbQKV_
+    """
+    @staticmethod
+    def forward(ctx, qkv, cos, sin, cos_k=None, sin_k=None):
+        """
+            qkv: (total, 3, nheads, headdim)
+            cos, sin: (seqlen, rotary_dim / 2)
+            cos_k, sin_k: (seqlen, rotary_dim / 2), optional
+        rotary_dim must be <= headdim
+        Apply rotary embedding *inplace* to the first rotary_dim of q and k.
+        """
+        _, three, _, headdim = qkv.shape
+        assert three == 3
+        rotary_seqlen, rotary_dim = cos.shape
+        rotary_dim *= 2
+        assert rotary_dim <= headdim
+        cos_k = cos if cos_k is None else cos_k
+        sin_k = sin if sin_k is None else sin_k
+        assert sin.shape == cos_k.shape == sin_k.shape == (rotary_seqlen,
+                                                           rotary_dim // 2)
+        q1, q2 = qkv[:, 0, :, :rotary_dim].chunk(2, dim=-1)
+        rotary_emb.apply_rotary(q1, q2, rearrange(cos, "s d -> s 1 d"),
+                                rearrange(sin, "s d -> s 1 d"), q1, q2, False)
+        k1, k2 = qkv[:, 1, :, :rotary_dim].chunk(2, dim=-1)
+        rotary_emb.apply_rotary(k1, k2, rearrange(cos_k, "s d -> s 1 d"),
+                                rearrange(sin_k, "s d -> s 1 d"), k1, k2,
+                                False)
+        ctx.save_for_backward(cos, sin, cos_k, sin_k)
+        return qkv
+
+    @staticmethod
+    def backward(ctx, dqkv):
+        cos, sin, cos_k, sin_k = ctx.saved_tensors
+        rotary_dim = cos.shape[-1]
+        rotary_dim *= 2
+        dq1, dq2 = dqkv[:, 0, :, :rotary_dim].chunk(2, dim=-1)
+        rotary_emb.apply_rotary(dq1, dq2, rearrange(cos, "s d -> s 1 d"),
+                                rearrange(sin, "s d -> s 1 d"), dq1, dq2, True)
+        dk1, dk2 = dqkv[:, 1, :, :rotary_dim].chunk(2, dim=-1)
+        rotary_emb.apply_rotary(dk1, dk2, rearrange(cos_k, "s d -> s 1 d"),
+                                rearrange(sin_k, "s d -> s 1 d"), dk1, dk2,
+                                True)
+        return dqkv, None, None, None, None
+
+
 class ConvertedInternLMRotaryEmbedding(torch.nn.Module):
     def __init__(self, dim: int, base=10000, scale_base=0, device=None):
         """ """
@@ -141,6 +200,23 @@ class ConvertedInternLMRotaryEmbedding(torch.nn.Module):
                 self._cos_k_cached = (torch.cos(freqs) / scale).to(x.dtype)
                 self._sin_k_cached = (torch.sin(freqs) / scale).to(x.dtype)
 
+    def forward(self,
+                qkv: torch.Tensor,
+                indexes=0) -> Tuple[torch.Tensor, torch.Tensor]:
+        self._update_cos_sin_cache(qkv, indexes)
+        if self.scale is None:
+            return apply_rotary_emb_qkv_(qkv, self._cos_cached[indexes],
+                                         self._sin_cached[indexes]).to(
+                                             qkv.dtype)
+        else:
+            return apply_rotary_emb_qkv_(
+                qkv,
+                self._cos_cached[indexes],
+                self._sin_cached[indexes],
+                self._cos_k_cached[indexes],
+                self._sin_k_cached[indexes],
+            ).to(qkv.dtype)
+
     def eval_forward(self, qkv, seqlen_offset=0):
         """
         seqlen_offset: can be used in generation where the qkv being passed in is only the last
@@ -161,6 +237,7 @@ class ConvertedInternLMRotaryEmbedding(torch.nn.Module):
             )
 
 
+apply_rotary_emb_qkv_ = ApplyRotaryEmbQKV_.apply
 legacy_apply_rotary_embed_qkv = LegacyApplyRotaryEmbQKV_.apply
 
 
@@ -1245,6 +1322,6 @@ class InternLMForCausalLM(InternLMPreTrainedModel):
         reordered_past = ()
         for layer_past in past_key_values:
             reordered_past += (tuple(
-                past_state.index_select(0, beam_idx)
+                past_state.index_select(0, beam_idx.to(past_state.device))
                 for past_state in layer_past), )
         return reordered_past

modeling_InternLM_XComposer.py

@@ -46,12 +46,13 @@ conversation
     def __init__(self, config):
         super().__init__(config)
 
-        print('Init VIT ... ', end='')
+        self.max_length = config.max_length
+        rank0_print('Init VIT ... ', end='')
         self.visual_encoder = create_eva_vit_g()
         self.ln_vision = LayerNorm(self.visual_encoder.num_features)
-        print('Done')
+        rank0_print('Done')
 
-        print('Init Perceive Sampler ... ', end='')
+        rank0_print('Init Perceive Sampler ... ', end='')
         with all_logging_disabled():
             self.Qformer, self.query_tokens = self.init_qformer(
                 config.num_query_token, self.visual_encoder.num_features)
@@ -61,9 +62,9 @@ conversation
                 layer.output = None
                 layer.intermediate = None
             self.Qformer.cls = None
-        print('Done')
+        rank0_print('Done')
 
-        print('Init InternLM ... ', end='')
+        rank0_print('Init InternLM ... ', end='')
         self.flag_image_start = nn.Parameter(torch.zeros([1, 1, 4096]))
         self.flag_image_end = nn.Parameter(torch.zeros([1, 1, 4096]))
         self.flag_image_start.requires_grad = False
@@ -81,7 +82,7 @@ conversation
             # speed up init llm
             with torch.device('meta'):
                 self.internlm_model = InternLMForCausalLM._from_config(config)
-            self.internlm_model.to_empty(device='cpu').to(torch.float16)
+            self.internlm_model.to_empty(device=config.device).to(torch.float16)
             self.internlm_model.to(config.device)
         for n, m in self.internlm_model.named_modules():
             if 'lora' in n:
@@ -89,7 +90,7 @@ conversation
 
         self.internlm_proj = nn.Linear(self.Qformer.config.hidden_size,
                                        self.internlm_model.config.hidden_size)
-        print('Done')
+        rank0_print('Done')
 
         self.vis_processor = transforms.Compose([
             transforms.Resize((224, 224),
@@ -111,6 +112,17 @@ conversation
             [StoppingCriteriaSub(stops=stop_words_ids)])
         self.gen_config['stopping_criteria'] = stopping_criteria
 
+        self.supports_gradient_checkpointing = True
+
+    def get_input_embeddings(self):
+        return self.internlm_model.get_input_embeddings()
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if value:
+            self.internlm_model.apply(
+                partial(self.internlm_model._set_gradient_checkpointing,
+                        value=True))
+
     def maybe_autocast(self, dtype=torch.float16):
         # if on cpu, don't use autocast
         # if on gpu, use autocast with dtype if provided, otherwise use torch.float16
@@ -268,3 +280,133 @@ conversation
         if history is not None:
             prompt_embeds = torch.cat([*history, prompt_embeds], dim=1)
         return prompt_embeds
+
+    ######################
+    #  code for training
+    ######################
+    def prompt_wrap(self, img_embeds, prompt):
+        batch_size = img_embeds.shape[0]
+        p_before, p_after = prompt.split('<ImageHere>')
+        p_before_tokens = self.tokenizer(p_before,
+                                         return_tensors="pt",
+                                         add_special_tokens=True).to(
+                                             img_embeds.device)
+
+        p_before_embeds = self.internlm_model.model.embed_tokens(
+            p_before_tokens.input_ids).expand(batch_size, -1, -1)
+        wrapped_img_embeds = torch.cat([p_before_embeds, img_embeds], dim=1)
+
+        wrapped_atts_img = torch.ones(wrapped_img_embeds.size()[:-1],
+                                      dtype=torch.long).to(img_embeds.device)
+
+        wrapped_target = torch.ones(
+            batch_size, wrapped_img_embeds.shape[1], dtype=torch.long).to(
+                img_embeds.device) * -100
+
+        return wrapped_img_embeds, wrapped_atts_img, wrapped_target
+
+    def align_text(self, samples, has_img=False):  ### add eos and eoa
+        text_new = []
+        if has_img:  ### remove the first user to wrap image features
+            text = [
+                t.replace("<image>", "").split("<|User|>:", 1)[-1].lstrip()
+                for t in samples["text_input"]
+            ]
+        else:
+            text = [t for t in samples["text_input"]]
+
+        text = [t + self.eoa + ' </s>' for t in text]
+        for i in range(len(text)):
+            temp = text[i]
+            temp = temp.replace('<|Bot|>', self.eoh + ' <|Bot|>')
+            temp = temp.replace(' <|User|>', self.eoa + ' <|User|>')
+            if temp.find(self.eoh) > temp.find(self.eoa):
+                temp = temp.replace(self.eoa, '', 1)
+            text_new.append(temp)
+        return text_new
+
+    def text2emb(self, text):
+        to_regress_tokens = self.tokenizer(text,
+                                           return_tensors="pt",
+                                           padding="longest",
+                                           truncation=True,
+                                           max_length=self.max_length,
+                                           add_special_tokens=False).to(
+                                               self.device)
+
+        targets = self.mask_human_targets(to_regress_tokens.input_ids)
+        targets = targets.to(self.device)
+
+        return to_regress_tokens, targets
+
+    def mask_human_targets(self, input_ids, pure=False):
+        target_batch = []
+        for bs in range(input_ids.shape[0]):
+            cur_idx = 0
+            ids = input_ids[bs]
+            targets = copy.deepcopy(ids)
+            last_eoa = 0
+            last_eoh = 0
+            for i, temp_id in enumerate(ids):
+                if temp_id == 103027:  #### end of human
+                    targets[cur_idx:i + 6] = -100
+                    cur_idx = i + 6
+                    last_eoh = i
+                elif temp_id == 103028:  ### end of assistant
+                    cur_idx = i + 1
+                    last_eoa = i
+                elif temp_id == 2:  ### eos and following pad
+                    targets[i + 1:] = -100  #### loss on eos, but not on pad
+                    break
+            if temp_id != 2 and last_eoa > last_eoh:  ### trunction, end at last question
+                targets[last_eoa +
+                        1:] = -100  #### mask all after the last answer
+
+            target_batch.append(targets.unsqueeze(0))
+
+        target_batch = torch.cat(target_batch, dim=0)
+        return target_batch
+
+    def forward(self,
+                input_ids=None,
+                attention_mask=None,
+                inputs_embeds=None,
+                labels=None,
+                output_attentions=None,
+                output_hidden_states=None,
+                return_dict=None,
+                **kwargs):
+
+        samples = kwargs.get('samples')
+        has_img = 'images' in samples.keys()
+
+        ### encode text
+        text = self.align_text(samples, has_img=has_img)
+        to_regress_tokens, targets = self.text2emb(text)
+
+        to_regress_embeds = self.internlm_model.model.embed_tokens(
+            to_regress_tokens.input_ids)
+        attention_mask = to_regress_tokens.attention_mask
+
+        if has_img:
+            header = samples["text_input"][0].split(' <|User|>:')[0]
+            prompt = header + ' <|User|>:<ImageHere>'
+
+            ### encode image
+            image = samples["image"]
+            img_embeds = self.encode_img(image)
+            img_embeds, atts_img, wrapped_target = self.prompt_wrap(
+                img_embeds, prompt)
+            ### combine text and image
+            to_regress_embeds = torch.cat([img_embeds, to_regress_embeds],
+                                          dim=1)
+            attention_mask = torch.cat([atts_img, attention_mask], dim=1)
+            targets = torch.cat([wrapped_target, targets], dim=1)
+
+        outputs = self.internlm_model(
+            inputs_embeds=to_regress_embeds,
+            attention_mask=attention_mask,
+            return_dict=True,
+            labels=targets,
+        )
+        return outputs

modeling_utils.py

@@ -2,12 +2,12 @@ import logging
 import math
 import os
 from contextlib import contextmanager
-from transformers import StoppingCriteria, StoppingCriteriaList
 
 import timm.models.hub as timm_hub
 import torch
 import torch.distributed as dist
 import torch.nn as nn
+from transformers import StoppingCriteria, StoppingCriteriaList
 
 
 def is_dist_avail_and_initialized():
@@ -28,12 +28,16 @@ def is_main_process():
     return get_rank() == 0
 
 
+def rank0_print(msg, **kwargs):
+    if is_main_process():
+        print(msg, **kwargs)
+
+
 def download_cached_file(url, check_hash=True, progress=False):
     """
     Download a file from a URL and cache it locally. If the file already exists, it is not downloaded again.
     If distributed, only the main process downloads the file, and the other processes wait for the file to be downloaded.
     """
-
     def get_cached_file_path():
         # a hack to sync the file path across processes
         parts = torch.hub.urlparse(url)
@@ -76,18 +80,16 @@ def all_logging_disabled(highest_level=logging.CRITICAL):
 
 
 class LoRALinear(nn.Linear):
-    def __init__(
-        self,
-        in_features: int,
-        out_features: int,
-        bias: bool = True,
-        device=None,
-        dtype=None,
-        lora_r=8,
-        lora_alpha=16,
-        lora_dropout=0.05,
-        **kwargs
-    ) -> None:
+    def __init__(self,
+                 in_features: int,
+                 out_features: int,
+                 bias: bool = True,
+                 device=None,
+                 dtype=None,
+                 lora_r=8,
+                 lora_alpha=16,
+                 lora_dropout=0.05,
+                 **kwargs) -> None:
         super().__init__(in_features, out_features, bias, device, dtype)
         self.lora_r = lora_r
         self.lora_alpha = lora_alpha
@@ -97,12 +99,16 @@ class LoRALinear(nn.Linear):
             self.lora_dropout = lambda x: x
         self.lora_scaling = self.lora_alpha / self.lora_r
 
-        self.lora_A = nn.Linear(
-            in_features, self.lora_r, bias=False, device=device, dtype=dtype
-        )
-        self.lora_B = nn.Linear(
-            self.lora_r, out_features, bias=False, device=device, dtype=dtype
-        )
+        self.lora_A = nn.Linear(in_features,
+                                self.lora_r,
+                                bias=False,
+                                device=device,
+                                dtype=dtype)
+        self.lora_B = nn.Linear(self.lora_r,
+                                out_features,
+                                bias=False,
+                                device=device,
+                                dtype=dtype)
 
         self.reset_parameters()
 
@@ -116,7 +122,8 @@ class LoRALinear(nn.Linear):
         orig_type = x.dtype
         res = super().forward(x)
         x = x.float()
-        res += self.lora_B(self.lora_A(self.lora_dropout(x))) * self.lora_scaling
+        res += self.lora_B(self.lora_A(
+            self.lora_dropout(x))) * self.lora_scaling
         return res.to(orig_type)
 
 
@@ -127,7 +134,7 @@ class StoppingCriteriaSub(StoppingCriteria):
 
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
         for stop in self.stops:
-            if torch.all((stop == input_ids[:, -len(stop) :])).item():
+            if torch.all((stop == input_ids[:, -len(stop):])).item():
                 return True
 
         return False