increase timeout and fix unit test

Signed-off-by: Yuki Huang <yukih@nvidia.com>

Author: yuki-97

tests/unit/models/generation/test_vllm_generation.py

@@ -1274,7 +1274,7 @@ async def test_vllm_refit_non_collocated_update_weights(
         print(f"Error during generation_cluster_separate shutdown: {e}")
 
 
-@pytest.mark.timeout(210)
+@pytest.mark.timeout(400)
 @pytest.mark.parametrize("tensor_parallel_size", [1, 2])
 def test_vllm_generation_with_megatron_training(
     cluster, tokenizer, tensor_parallel_size