GRPO最后一个step训练完毕后触发AssertionError

[fwensen]

Checklist / 检查清单

• I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

使用example/megatron/grpo/dence_colocate.sh demo时报错如下

lsg1a-reservedpool-ljhxg-46:310402:312269 [0] NCCL INFO [Service thread] Connection closed by localRank 3
lsg1a-reservedpool-ljhxg-46:310404:311949 [2] NCCL INFO [Service thread] Connection closed by localRank 3
lsg1a-reservedpool-ljhxg-46:310402:311951 [0] NCCL INFO [Service thread] Connection closed by localRank 3
[rank7]: Traceback (most recent call last):
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/cli/_megatron/rlhf.py", line 24, in <module>
[rank7]:     megatron_rlhf_main()
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/megatron/pipelines/train/rlhf.py", line 73, in megatron_rlhf_main
[rank7]:     return MegatronRLHF(args).main()
[rank7]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/pipelines/base.py", line 52, in main
[rank7]:     result = self.run()
[rank7]:              ^^^^^^^^^^
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/megatron/pipelines/train/sft.py", line 68, in run
[rank7]:     trainer.train(train_dataset, val_dataset)
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/megatron/trainers/grpo_trainer.py", line 70, in train
[rank7]:     super().train(train_dataset, val_dataset)
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/megatron/trainers/base.py", line 739, in train
[rank7]:     self.run_train_step(train_data_iterator, val_data_iterator)
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/megatron/trainers/base.py", line 679, in run_train_step
[rank7]:     metrics, grad_norm, update_successful = self.train_step(train_data_iterator)
[rank7]:                                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/megatron/trainers/base.py", line 922, in train_step
[rank7]:     data_iterator = self._replace_data_iterator(train_data_iterator)
[rank7]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/megatron/trainers/rollout_mixin.py", line 164, in _replace_data_ite
rator
[rank7]:     self._buffered_inputs = self._build_rollout_buffer(data_iterator)
[rank7]:                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/megatron/trainers/grpo_trainer.py", line 234, in _build_rollout_buf
fer
[rank7]:     micro_batch_data = self._generate_and_score_completions(rollout_batch)
[rank7]:                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/megatron/trainers/grpo_trainer.py", line 258, in _generate_and_scor
e_completions
[rank7]:     samples = self._generate_completions(samples)
[rank7]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/rlhf_trainers/utils.py", line 594, in wrapper
[rank7]:     return func(self, *args, **kwargs)
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/megatron/trainers/rollout_mixin.py", line 751, in _generate_complet
ions
[rank7]:     outputs: List[RolloutOutput] = self._rollout(samples)
[rank7]:                                    ^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/megatron/trainers/rollout_mixin.py", line 785, in _rollout
[rank7]:     return self._colocate_rollout(samples, request_config)
[rank7]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/megatron/trainers/rollout_mixin.py", line 852, in _colocate_rollout
[rank7]:     outputs: List[RolloutOutput] = self.engine.infer(
[rank7]:                                    ^^^^^^^^^^^^^^^^^^
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/infer_engine/grpo_vllm_engine.py", line 41, in infer
[rank7]:     res = super().infer(
[rank7]:           ^^^^^^^^^^^^^^
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/infer_engine/vllm_engine.py", line 853, in infer
[rank7]:     self._add_request(inputs, generation_config, request_id, adapter_request=adapter_request)
[rank7]:   File "/mnt/tidalfs-alsg01/dataset/rednote/shenya/code/ms-swift/swift/infer_engine/vllm_engine.py", line 496, in _add_request
[rank7]:     return self.engine.add_request(request_id, llm_inputs, generation_config, **kwargs)
[rank7]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/usr/local/lib/python3.12/site-packages/vllm/v1/engine/llm_engine.py", line 274, in add_request
[rank7]:     self.engine_core.add_request(request)
[rank7]:   File "/usr/local/lib/python3.12/site-packages/vllm/v1/engine/core_client.py", line 296, in add_request
[rank7]:     req, request_wave = self.engine_core.preprocess_add_request(request)
[rank7]:                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/usr/local/lib/python3.12/site-packages/vllm/v1/engine/core.py", line 747, in preprocess_add_request
[rank7]:     request.mm_features = self.mm_receiver_cache.get_and_update_features(
[rank7]:                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/usr/local/lib/python3.12/site-packages/vllm/multimodal/cache.py", line 591, in get_and_update_features
[rank7]:     feature.data = self.get_and_update_item(feature.data, cache_key)
[rank7]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/usr/local/lib/python3.12/site-packages/vllm/multimodal/cache.py", line 644, in get_and_update_item
[rank7]:     assert mm_item is not None, f"Expected a cached item for {mm_hash=}"
[rank7]:            ^^^^^^^^^^^^^^^^^^^
[rank7]: AssertionError: Expected a cached item for mm_hash='6813bfe0cb9728171ee26240e8ad9b106aaac6fdba9a5181a7a04eaae43952ff'

How to Reproduce / 如何复现

example/megatron/grpo/dence_colocate.sh

Additional Information / 补充信息

No response

[hjh0119]

it's a known issue in vllm

plz set vllm_mm_processor_cache_gb 0