model training with normal image

[egbertYeah]

In the src/PaDT/trainer/padt_sft_trainer.py file, the compute_loss function does not properly handle normal images (those without target objects). When the training data includes items labeled as "No objects from the list are present in the image," the program encounters an error.

[Dyb3438]

@egbertYeah Could you give the error message?

[egbertYeah]

same as i said, i used instances_val2017.json for simple ovd task, when samples include normal image, an error apper as below pic.

@Dyb3438

[egbertYeah]

i rewrite compute_loss funtion to support normal sample, the error still appear.

def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
    if return_outputs:
        raise ValueError("The PaDTTrainer does not support returning outputs")

    # 1. 预处理提示文本
    prompt_text = [self.processing_class.apply_chat_template(i['prompt'], tokenize=False, add_generation_prompt=True) for i in inputs]
    
    images = []
    completions = []
    solutions = []
    
    # 2. 处理图像和生成补全文本
    has_objects_in_batch = False  # 标记批次中是否有对象
    
    for idx, x in enumerate(inputs):
        # 处理图像
        assert len(x['image_path']) == 1, "current support only an image per sample"
        
        for img in x['image_path']:
            image = PIL.Image.open(img)
            try:
                w, h = image.size
                if w < 28 or h < 28:
                    if w < h:
                        new_w = 28
                        new_h = int(h * (28 / w))
                    else:
                        new_h = 28
                        new_w = int(w * (28 / h))
                    image = image.resize((new_w, new_h), PIL.Image.Resampling.LANCZOS)
            except:
                pass
            images.append(image)
        
        # 处理补全文本
        im_w, im_h = image.size
        patch_w, patch_h = round(im_w / 28), round(im_h / 28)

        solution = x['solution']
        
        # 检查是否有对象
        if 'objects' in solution and len(solution['objects']) > 0:
            has_objects_in_batch = True
            '''
            >>> completion = "In this image, there are 11 \"person\" (<|Obj_0|>, <|Obj_1|>, <|Obj_2|>, <|Obj_3|>, <|Obj_4|>, <|Obj_5|>, <|Obj_6|>, <|Obj_7|>, <|Obj_8|>, <|Obj_9|>, <|Obj_10|>), 1 \"cup\" (<|Obj_11|>), 1 \"bowl\" (<|Obj_12|>), 1 \"couch\" (<|Obj_13|>), 1 \"cell phone\" (<|Obj_14|>)."

            >>> obj_strs
            ['<|Obj_0|>', '<|Obj_1|>', '<|Obj_2|>', '<|Obj_3|>', '<|Obj_4|>', '<|Obj_5|>', '<|Obj_6|>', '<|Obj_7|>', '<|Obj_8|>', '<|Obj_9|>', '<|Obj_10|>', '<|Obj_11|>', '<|Obj_12|>', '<|Obj_13|>', '<|Obj_14|>']

            >>> completion_parts
            ['In this image, there are 11 "person" (', ', ', ', ', ', ', ', ', ', ', ', ', ', ', ', ', ', ', ', ', '), 1 "cup" (', '), 1 "bowl" (', '), 1 "couch" (', '), 1 "cell phone" (', ').']
            '''
            # 有对象的情况：正常处理VRT
            completion = solution['text']
            pattern = r'(<\|Obj_(\d+)\|>)'
            obj_in_completion = re.findall(pattern, completion)
            obj_strs = [i[0] for i in obj_in_completion]
            objs = [solution['objects'][int(i[1])] for i in obj_in_completion]
            pattern_without_matching = r'<\|Obj_\d+\|>'
            completion_parts = re.split(pattern_without_matching, completion)
            
            completion_with_vrt = completion_parts[0]
            new_objs = []
            
            for obj_str, completion_part, obj in zip(obj_strs, completion_parts[1:], objs):
                obj_ = obj.copy()
                selected_patches = np.array(obj_['patches'])
                
                # Patch选择策略
                if self.args.random_select_patch_num < 0:
                    pick_patch = selected_patches.copy()
                elif self.args.random_select_patch is False:
                    # 策略性选择：中心、左、上、右、下
                    selected_patches_x = selected_patches % patch_w
                    selected_patches_y = selected_patches // patch_w
                    
                    left_patches_m = selected_patches_x == selected_patches_x.min()
                    right_patches_m = selected_patches_x == selected_patches_x.max()
                    top_patches_m = selected_patches_y == selected_patches_y.min()
                    bottom_patches_m = selected_patches_y == selected_patches_y.max()
                    centre_patches_m = (left_patches_m + right_patches_m + top_patches_m + bottom_patches_m) == 0                        
                    
                    left_patches = selected_patches[left_patches_m]
                    right_patches = selected_patches[right_patches_m]
                    top_patches = selected_patches[top_patches_m]
                    bottom_patches = selected_patches[bottom_patches_m]
                    centre_patches = selected_patches[centre_patches_m]
                    
                    if centre_patches_m.sum() == 0:
                        centre_patches = selected_patches
                        
                    pick_patch = np.array([
                        np.random.choice(centre_patches), 
                        np.random.choice(left_patches), 
                        np.random.choice(top_patches), 
                        np.random.choice(right_patches), 
                        np.random.choice(bottom_patches)
                    ])
                else:
                    # 随机选择
                    if selected_patches.shape[0] < self.args.random_select_patch_num:
                        pick_patch = np.random.choice(selected_patches, self.args.random_select_patch_num, replace=True)
                    else:
                        pick_patch = np.random.choice(selected_patches, self.args.random_select_patch_num, replace=False)
                
                obj_['picked'] = pick_patch
                new_objs.append(obj_)
                completion_with_vrt += self.processing_class.pid2vrt(pick_patch) + completion_part

            solutions.append({
                'text': solution['text'],
                'objects': new_objs
            })
            # 构建带有VRT的补全文本
            completions.append(completion_with_vrt + self.processing_class.tokenizer.eos_token)
            
        else:
            # 没有对象的情况：使用原始文本，不添加VRT
            completion = solution['text']
            solutions.append({
                'text': solution['text'],
                'objects': []  # 明确标记为空
            })
            completions.append(completion + self.processing_class.tokenizer.eos_token)

    # 3. Tokenizing
    prompt_inputs = self.processing_class(
        text=prompt_text,
        images=images,
        return_tensors='pt',
        padding=True,
        padding_side='left',
        add_special_tokens=False
    )
    prompt_inputs = super()._prepare_inputs(prompt_inputs)
    prompt_ids, prompt_mask = prompt_inputs["input_ids"], prompt_inputs["attention_mask"]
    batch_size = prompt_ids.size(0)
    prompt_length = prompt_ids.size(1)
    multimodal_inputs = {
        'image_grid_thw': prompt_inputs['image_grid_thw'],
        'pixel_values': prompt_inputs['pixel_values']
    }
    
    completion_inputs = self.processing_class(
        text=completions,
        return_tensors='pt',
        padding=True,
        padding_side='right',
        add_special_tokens=False
    )
    completion_inputs = super()._prepare_inputs(completion_inputs)
    completion_ids, completion_mask = completion_inputs["input_ids"], completion_inputs["attention_mask"]

    # 4. 构建损失掩码 - 处理正常图像的情况
    loss_masks = []
    gt_bboxes = []
    
    vision_patch_nums = torch.nn.functional.pad(
        (multimodal_inputs['image_grid_thw'].cumprod(-1)[:, -1] // 
        (self.model.config.vision_config.spatial_merge_size ** 2)).cumsum(-1), 
        (1, 0), 'constant', 0
    )
    all_vision_patch_nums = vision_patch_nums[-1]
    
    # 检查是否有对象需要构建损失掩码
    objects_found = False
    
    for sol, vpn in zip(solutions, vision_patch_nums[:-1]):
        if 'objects' in sol and len(sol['objects']) > 0:
            for obj in sol['objects']:
                # 检查picked patches是否有效
                if ('picked' not in obj or not hasattr(obj['picked'], 'shape') or 
                    len(obj['picked'].shape) == 0 or obj['picked'].shape[0] == 0):
                    continue
                    
                objects_found = True
                this_object_loss_mask = torch.zeros(
                    (obj['picked'].shape[0], all_vision_patch_nums), 
                    device=self.accelerator.device, dtype=torch.bool
                )
                
                # 标记相关patches
                if 'patches' in obj and len(obj['patches']) > 0:
                    this_object_loss_mask[:, vpn.item() + np.array(obj['patches'])] = True
                
                # 排除选中的patches
                this_object_loss_mask[np.arange(obj['picked'].shape[0]), vpn.item() + obj['picked']] = False
                loss_masks.append(this_object_loss_mask)
                gt_bboxes.append(obj['bbox'])

    # 处理空损失掩码的情况（正常图像）
    # if len(loss_masks) == 0:
    #     # 创建空的损失掩码
    #     loss_masks = torch.zeros((0, all_vision_patch_nums), device=self.accelerator.device, dtype=torch.bool)
    # else:
    #     loss_masks = torch.cat(loss_masks, dim=0)
    # 关键修复：确保所有GPU上的张量形状一致
    if len(loss_masks) == 0:
        # 创建形状一致的空白张量
        print("WARNING: loss_masks is empty! Creating dummy tensors...")
        # 在代码关键位置添加打印语句
        print(f"Number of solutions: {len(solutions)}")
        print(f"vision_patch_nums: {vision_patch_nums}")
        print(f"vision_patch_nums[:-1]: {vision_patch_nums[:-1]}")

        print(f"loss_masks length: {len(loss_masks)}")
        print(f"image path: {[x['image_path'] for x in inputs]}")

        loss_masks = torch.zeros((0, all_vision_patch_nums), device=self.accelerator.device, dtype=torch.bool)
        gt_bboxes = torch.zeros((0, 4), device=self.accelerator.device, dtype=torch.bfloat16)
    else:
        loss_masks = torch.cat(loss_masks, dim=0)
        gt_bboxes = torch.tensor(gt_bboxes, device=self.accelerator.device, dtype=torch.bfloat16)
        if len(gt_bboxes.shape) == 1:
            gt_bboxes = gt_bboxes.unsqueeze(dim=-1).repeat_interleave(4, dim=-1)

    # 填充文本部分
    loss_masks = torch.nn.functional.pad(loss_masks, (self.model_embed_token_size, 0), 'constant', False)
    
    # 处理边界框
    # if len(gt_bboxes) > 0:
    #     gt_bboxes = torch.Tensor(gt_bboxes).to(self.accelerator.device).to(torch.bfloat16)
    #     if len(gt_bboxes.shape) == 1:
    #         gt_bboxes = gt_bboxes.unsqueeze(dim=-1).repeat_interleave(4, dim=-1)
    # else:
    #     # 没有边界框的情况
    #     gt_bboxes = torch.zeros((0, 4), device=self.accelerator.device, dtype=torch.bfloat16)

    # 5. 模型前向传播
    input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
    attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
    input_ids = self.processing_class.assign_to_global_vrt_id(input_ids, multimodal_inputs['image_grid_thw'])

    model_output = model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True, **multimodal_inputs)

    # 6. 多任务损失计算
    logits = model_output.logits[:, prompt_length-1:-1, :]  # shape: (batch_size, seq_len, vocab_size)
    input_ids = input_ids[:, prompt_length:]        # 模型输出对应的ID
    if len(loss_masks) == 0: print("WARNING: loss_masks is empty before visual patch masking!")
    # 只在有对象时应用视觉patch掩码
    if self.args.use_sft_vp_mask and objects_found and loss_masks.shape[0] > 0:
        visual_patch_mask = input_ids >= self.model_embed_token_size
        # 确保loss_masks与visual_patch_mask形状兼容
        if loss_masks.shape[0] > 0:
            logits[visual_patch_mask] = logits[visual_patch_mask].masked_fill(loss_masks, float('-inf'))
    
    # 文本生成损失（所有样本都有）
    logit_log_probs = logits.log_softmax(dim=-1)
    token_log_prob = torch.gather(logit_log_probs, dim=-1, index=input_ids.unsqueeze(-1)).squeeze(-1)
    per_token_loss = -token_log_prob
    sft_loss = ((per_token_loss * completion_mask).sum(dim=-1) / (completion_mask.sum(dim=-1) + 1e-4)).to(logits.dtype)
    # self._metrics['sft_loss'].append(self.accelerator.gather_for_metrics(sft_loss).mean().item())
    # 使用accelerator.gather确保分布式一致性
    gathered_sft_loss = self.accelerator.gather_for_metrics(sft_loss)
    self._metrics['sft_loss'].append(gathered_sft_loss.mean().item() if gathered_sft_loss.numel() > 0 else 0.0)
    if len(loss_masks) == 0: print("WARNING: loss_masks is empty after text generation loss computation!")
    # 7. 视觉任务损失 - 关键修复：确保分布式一致性
    bbox_loss = torch.tensor(0.0, device=self.accelerator.device)
    mask_loss = torch.tensor(0.0, device=self.accelerator.device)
    score_loss = torch.tensor(0.0, device=self.accelerator.device)
    if objects_found and gt_bboxes.shape[0] > 0:
        try:
            # 解码视觉输出
            hidden_states = torch.stack(model_output.hidden_states, dim=1)[:, -1:, prompt_length-1:-1].permute(2, 1, 0, 3).unsqueeze(dim=-2).contiguous()
            completions, feats, labels, vps, vps_feats = parseVRTintoCompletion(
                self.processing_class, completion_ids, hidden_states, 
                torch.tensor([False] * batch_size), model_output.past_image_embeds, 
                multimodal_inputs['image_grid_thw']
            )
            
            low_res_image_embeds = model_output.past_image_embeds
            high_res_image_embeds = model_output.past_high_res_image_embeds
            visual_pe = model_output.past_visual_pe
            
            # 热身阶段
            if self.state.epoch < (self.state.num_train_epochs / 4) and self.state.global_step < 300 and self.args.use_warm_up:
                feats = vps_feats
                
            decoded_list = model(feats, low_res_image_embeds, high_res_image_embeds, multimodal_inputs['image_grid_thw'], visual_pe, is_main=False)
            
            # 边界框损失
            if self.args.use_bbox_loss and gt_bboxes.shape[0] > 0:
                pred_bboxes = decoded_list['pred_boxes']
                num_bboxes = gt_bboxes.shape[0]
                
                # 确保所有GPU上的bbox数量一致
                giou, iou = self.generalized_box_iou(self.box_cxcywh_to_xyxy(pred_bboxes), gt_bboxes)
                giou = torch.diag(giou).to(pred_bboxes.dtype)
                bbox_loss = 1. - giou.sum() / (num_bboxes + 1e-4)
                bbox_loss += torch.nn.functional.l1_loss(pred_bboxes, self.box_xyxy_to_cxcywh(gt_bboxes), reduction='none').sum() / (num_bboxes + 1e-4)
                
                # 使用gather确保分布式一致性
                gathered_bbox_loss = self.accelerator.gather_for_metrics(bbox_loss.unsqueeze(0))
                self._metrics['bbox_loss'].append(gathered_bbox_loss.mean().item() if gathered_bbox_loss.numel() > 0 else 0.0)
                
                gathered_iou = self.accelerator.gather_for_metrics(torch.diag(iou).sum().unsqueeze(0) / (num_bboxes + 1e-4))
                self._metrics['iou'].append(gathered_iou.mean().item() if gathered_iou.numel() > 0 else 0.0)
                
                gathered_giou = self.accelerator.gather_for_metrics(giou.sum().unsqueeze(0) / (num_bboxes + 1e-4))
                self._metrics['giou'].append(gathered_giou.mean().item() if gathered_giou.numel() > 0 else 0.0)

            # 掩码损失和分数损失类似处理...
            
        except Exception as e:
            print(f"Visual task loss computation failed: {e}")
            # 在错误情况下设置默认值，避免分布式超时
            bbox_loss = torch.tensor(0.0, device=self.accelerator.device)
            mask_loss = torch.tensor(0.0, device=self.accelerator.device)
            score_loss = torch.tensor(0.0, device=self.accelerator.device)
    else:
        # 没有对象的情况，设置默认指标值
        self._metrics['bbox_loss'].append(0.0)
        self._metrics['mask_loss'].append(0.0)
        self._metrics['score_loss'].append(0.0)
        self._metrics['iou'].append(0.0)
        self._metrics['giou'].append(0.0)
    if len(loss_masks) == 0: print("WARNING: loss_masks is empty before final loss computation!")
    # 7. 总损失
    loss = sft_loss.mean() + bbox_loss + score_loss + mask_loss
    if len(loss_masks) == 0: print(f"WARNING: loss_masks is empty after final loss computation! Loss = {loss.item()}")
    # 记录正常图像的比例
    # normal_image_ratio = 1.0 - (float(objects_found) / len(inputs)) if len(inputs) > 0 else 0.0
    # self._metrics['normal_image_ratio'].append(normal_image_ratio)
    
    return loss

[egbertYeah]

@Dyb3438

[Dyb3438]

@egbert Yeah, I understand your issue. It is caused by different behaviors on different GPU cards. For example, GPU 0 forwards some VRT features to the decoder, so it saves activations in the decoder module. However, GPU 1 has samples without VRTs, so no activations are saved. As a result, during backpropagation, the decoder states become inconsistent across GPUs, which causes the training process to hang due to synchronization failure.

[Dyb3438]

In our current experiments, we use a larger batch size (e.g., 16 samples per GPU), so the probability that all 16 images contain no target simultaneously is extremely low — almost zero. Therefore, this issue didn’t appear during our experiments and was not noticed at first.

[Dyb3438]

@egbertYeah I will fix this issue after my recent rebuttal stage.

[egbertYeah]

@egbert Yeah, I understand your issue. It is caused by different behaviors on different GPU cards. For example, GPU 0 forwards some VRT features to the decoder, so it saves activations in the decoder module. However, GPU 1 has samples without VRTs, so no activations are saved. As a result, during backpropagation, the decoder states become inconsistent across GPUs, which causes the training process to hang due to synchronization failure.

Your answer is very insightful. So, how can this issue be avoided? I am using an H100 with a batch size of 2.

[egbertYeah]

@egbertYeah I will fix this issue after my recent rebuttal stage.

Can you give me some suggestions ？

[Dyb3438]

@egbert Yeah, if you prefer not to modify the source code, there are several workarounds you can try:

1. Run the experiment on a single GPU only, so that all samples share the same decoder behavior.

2. Increase the batch size as much as possible, which reduces the chance of having no VRT samples on a specific GPU.

3. Modify the sampling strategy in the dataloader to ensure that samples with VRTs appear on every GPU simultaneously.

[egbertYeah]

@egbert Yeah, if you prefer not to modify the source code, there are several workarounds you can try:

1. Run the experiment on a single GPU only, so that all samples share the same decoder behavior.
2. Increase the batch size as much as possible, which reduces the chance of having no VRT samples on a specific GPU.
3. Modify the sampling strategy in the dataloader to ensure that samples with VRTs appear on every GPU simultaneously.

@Dyb3438 thanks，i will try your advices.

[egbertYeah]

@Dyb3438 i modify the sampler class, it can solve this promblem.

class RepeatRandomSampler(Sampler):
"""
Sampler that repeats the indices of a dataset in a structured manner.

Args:
    data_source (`Sized`):
        Dataset to sample from.
    mini_repeat_count (`int`):
        Number of times to repeat each index per batch.
    batch_size (`int`, *optional*, defaults to `1`):
        Number of unique indices per batch.
    repeat_count (`int`, *optional*, defaults to `1`):
        Number of times to repeat the full sampling process.
    seed (`int` or `None`, *optional*, defaults to `None`):
        Random seed for reproducibility.
    max_retry (`int`, *optional*, defaults to `10`):
        Maximum number of retries when a mini-batch contains only normal images.
"""

def __init__(
    self,
    data_source: Sized,
    mini_repeat_count: int,
    batch_size: int = 1,
    repeat_count: int = 1,
    seed: Optional[int] = None,
    num_processes: int = 1,
    gradient_accumulation_steps: int = 1,
    per_device_batch_size: int = 1,
    max_retry: int = 20
):
    self.data_source = data_source
    self.mini_repeat_count = mini_repeat_count
    self.batch_size = batch_size
    self.repeat_count = repeat_count
    self.num_samples = len(data_source)
    self.seed = seed
    self.num_processes = num_processes
    self.gradient_accumulation_steps = gradient_accumulation_steps
    self.max_retry = max_retry
    self.generator = torch.Generator()
    self.per_device_batch_size = per_device_batch_size
    if seed is not None:
        self.generator.manual_seed(seed)

def _has_objects(self, idx):
    """检查样本是否包含对象"""
    try:
        item = self.data_source[idx]
        if ('solution' in item and 
            'objects' in item['solution'] and 
            len(item['solution']['objects']) > 0):
            return True
        return False
    except:
        return False

def _is_valid_mini_batch(self, chunk):
    """检查mini-batch是否包含至少一个有对象的样本"""
    # print(f"chunk: {chunk}")
    for idx in range(0, len(chunk), self.per_device_batch_size):
        for j in range(idx, idx+self.per_device_batch_size):
            if not self._has_objects(chunk[j]):
                # print(f"error. {chunk}")
                return False
    return True

def __iter__(self):
    indexes = torch.randperm(self.num_samples, generator=self.generator).tolist()
    indexes = [indexes[i : i + self.batch_size] for i in range(0, len(indexes) // self.batch_size * self.batch_size, self.batch_size)]
    
    for chunk in indexes:
        for _ in range(self.repeat_count):
            for accumultaion_step_idx in range(self.gradient_accumulation_steps):
                accumulation_chunk = chunk[accumultaion_step_idx::self.gradient_accumulation_steps]
                
                # 检查当前mini-batch是否都是正常图
                retry_count = 0
                while not self._is_valid_mini_batch(accumulation_chunk) and retry_count < self.max_retry:
                    # 重新采样这个mini-batch
                    retry_count += 1
                    # 从整个数据集中随机选择新的样本替换当前mini-batch
                    new_chunk = torch.randperm(self.num_samples, generator=self.generator).tolist()[:len(accumulation_chunk)]
                    accumulation_chunk = new_chunk
                
                if retry_count >= self.max_retry:
                    print(f"Warning: Failed to find a valid mini-batch after {self.max_retry} retries")
                
                chunk_idx = torch.arange(0, len(accumulation_chunk), dtype=torch.long)[:, None].repeat_interleave(self.mini_repeat_count, -1).view(-1)
                for index in chunk_idx:
                    yield accumulation_chunk[index]

def __len__(self) -> int:
    return self.num_samples * self.mini_repeat_count * self.repeat_count