[bugfix] fix MiniCPMV4_6 text-only batch graph break in distributed training#9443
Open
randydl wants to merge 3 commits into
Open
[bugfix] fix MiniCPMV4_6 text-only batch graph break in distributed training#9443randydl wants to merge 3 commits into
randydl wants to merge 3 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a _post_encode method in swift/template/templates/minicpm.py to handle text-only batches during training by passing a dummy tensor through the vision encoder, ensuring consistent parameter usage across ranks. The reviewer suggested several defensive programming improvements to prevent runtime errors, including unwrapping DDP/FSDP model wrappers, retrieving the model's data type from inputs_embeds instead of model.dtype, and safely handling pooler_output types before concatenation.
Comment on lines
+690
to
+723
| def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]: | ||
| if not self.is_training: | ||
| return inputs | ||
|
|
||
| pixel_values = inputs.get('pixel_values') | ||
| pixel_values_videos = inputs.get('pixel_values_videos') | ||
|
|
||
| if pixel_values is None and pixel_values_videos is None: | ||
| # Text-only batch: run a minimal dummy through the vision encoder | ||
| # so every rank's computation graph includes the same parameters. | ||
| device = inputs['input_ids'].device | ||
| patch_size = model.config.vision_config.patch_size | ||
|
|
||
| dummy_pv = torch.zeros( | ||
| 1, 3, 4 * patch_size, 4 * patch_size, | ||
| device=device, dtype=model.dtype | ||
| ) | ||
| dummy_ts = torch.tensor( | ||
| [[4, 4]], device=device, dtype=torch.int32 | ||
| ) | ||
|
|
||
| inputs_embeds = model.get_input_embeddings()(inputs['input_ids']) | ||
| vision_output = model.get_image_features( | ||
| dummy_pv, dummy_ts, downsample_mode=self.downsample_mode | ||
| ) | ||
|
|
||
| dummy_feats = torch.cat(vision_output.pooler_output, dim=0).to( | ||
| device=inputs_embeds.device, dtype=inputs_embeds.dtype | ||
| ) | ||
| inputs_embeds = inputs_embeds + dummy_feats.mean() * 0.0 | ||
|
|
||
| return {'inputs_embeds': inputs_embeds} | ||
|
|
||
| return inputs |
Contributor
There was a problem hiding this comment.
To ensure robustness and prevent potential runtime errors during training, we should address a few defensive programming concerns:
- DDP/FSDP Unwrapping: If the model is wrapped in
DistributedDataParallel(DDP) or FSDP, accessing attributes likeconfigor methods likeget_input_embeddingsdirectly onmodelwill raise anAttributeErrorbecause these wrappers do not delegate attribute access. Unwrapping the inner module viamodel.moduleif present is highly recommended. - Model Dtype:
model.dtypeis not a standard PyTorchnn.Moduleattribute and might not be available on all custom wrappers or PEFT models. Getting the dtype directly frominputs_embeds.dtypeis much safer. - Defensive handling of
pooler_output: Depending on the model configuration or future updates,vision_output.pooler_outputmight be a single tensor instead of a list of tensors. We should check its type before callingtorch.catto avoid potential type errors.
def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
if not self.is_training:
return inputs
pixel_values = inputs.get('pixel_values')
pixel_values_videos = inputs.get('pixel_values_videos')
if pixel_values is None and pixel_values_videos is None:
# Unwrap DDP/FSDP wrapper if present
if hasattr(model, 'module'):
model = model.module
# Text-only batch: run a minimal dummy through the vision encoder
# so every rank's computation graph includes the same parameters.
inputs_embeds = model.get_input_embeddings()(inputs['input_ids'])
device = inputs_embeds.device
dtype = inputs_embeds.dtype
patch_size = model.config.vision_config.patch_size
dummy_pv = torch.zeros(
1, 3, 4 * patch_size, 4 * patch_size,
device=device, dtype=dtype
)
dummy_ts = torch.tensor(
[[4, 4]], device=device, dtype=torch.int32
)
vision_output = model.get_image_features(
dummy_pv, dummy_ts, downsample_mode=self.downsample_mode
)
pooler_output = vision_output.pooler_output
if isinstance(pooler_output, (list, tuple)):
dummy_feats = torch.cat(pooler_output, dim=0)
else:
dummy_feats = pooler_output
dummy_feats = dummy_feats.to(device=device, dtype=dtype)
inputs_embeds = inputs_embeds + dummy_feats.mean() * 0.0
return {'inputs_embeds': inputs_embeds}
return inputs
Contributor
Author
There was a problem hiding this comment.
_post_encode 收到的 model 已经是 unwrap 后的，不需要再处理。hook 注册时经过 accelerator.unwrap_model(model) 拿到原始模型，然后 register_forward_pre_hook 注册在这个对象上。训练时 wrapper 调用 self.module() 触发的是原始模型上的 hook，所以回调中的 model 就是 MiniCPMV4_6ForConditionalGeneration 本身，直接访问 .model.language_model.embed_tokens 没问题。
randydl
changed the title
added 2 commits
May 29, 2026 01:36
Contributor
Author
Before Fix:After Fix: |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Added a _post_encode override in MiniCPMV4_6Template that intercepts text-only batches and injects a zero-contribution dummy forward pass through the vision encoder: