[AMD] rocm 7.2 image release, PR test, Nightly Test (#17799)

Co-authored-by: Alan Kao <akao@amd.com>
Co-authored-by: bingxche <Bingxu.Chen@amd.com>
Co-authored-by: Michael <13900043+michaelzhang-ai@users.noreply.github.com>

test/registered/rl/test_update_weights_from_distributed.py

@@ -37,6 +37,7 @@ from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
+    is_in_amd_ci,
     is_in_ci,
     popen_launch_server,
 )
@@ -64,6 +65,60 @@ def verify_params_not_close(params1, params2, error_msg):
     assert not np.allclose(np.array(params1), np.array(params2)), error_msg
 
 
+def _warmup_broadcast(
+    hf_base_model,
+    state_dict_key_to_shape,
+    tie_word_embeddings,
+    load_format,
+    group,
+):
+    """Run one broadcast round to warm up RCCL before timing."""
+    broadcast_parameters = list(state_dict_key_to_shape.keys())
+    if tie_word_embeddings:
+        broadcast_parameters.remove("lm_head.weight")
+
+    if load_format == "flattened_bucket":
+        named_tensors = [
+            (name, hf_base_model.get_parameter(name)) for name in broadcast_parameters
+        ]
+        bucket = FlattenedTensorBucket(named_tensors=named_tensors)
+        flattened_tensor = bucket.get_flattened_tensor()
+        torch.distributed.broadcast(flattened_tensor, src=0, group=group)
+    else:
+        for name in broadcast_parameters:
+            torch.distributed.broadcast(
+                hf_base_model.get_parameter(name),
+                src=0,
+                group=group,
+            )
+
+
+def _warmup_update(
+    backend, engine, url, names, dtypes, shapes, load_format, pause_generation_mode
+):
+    """Run one update round to warm up RCCL before timing."""
+    if backend == "Engine":
+        engine.update_weights_from_distributed(
+            names,
+            dtypes=dtypes,
+            shapes=shapes,
+            group_name="test_parameter_update_group",
+            load_format=load_format,
+        )
+    else:
+        requests.post(
+            f"{url}/update_weights_from_distributed",
+            json={
+                "names": names,
+                "dtypes": dtypes,
+                "shapes": shapes,
+                "group_name": "test_parameter_update_group",
+                "load_format": load_format,
+                "flush_cache": not (pause_generation_mode == "in_place"),
+            },
+        )
+
+
 def init_process(
     rank,
     world_size,
@@ -180,6 +235,18 @@ def init_process_hf(
     )
     torch.cuda.synchronize()
     barrier.wait()
+
+    # Warmup: trigger RCCL initialization so it's excluded from timing
+    if is_in_amd_ci():
+        _warmup_broadcast(
+            hf_base_model,
+            state_dict_key_to_shape,
+            tie_word_embeddings,
+            load_format,
+            group,
+        )
+        torch.cuda.synchronize()
+
     time_begin_broadcast = time.perf_counter()
 
     # The last parameter is lm_head.weight, which is tied
@@ -354,6 +421,21 @@ def init_process_sgl(
         )
     torch.cuda.synchronize()
     barrier.wait()
+
+    # Warmup: trigger RCCL initialization so it's excluded from timing
+    if is_in_amd_ci():
+        _warmup_update(
+            backend,
+            engine if backend == "Engine" else None,
+            url if backend != "Engine" else None,
+            names,
+            dtypes,
+            shapes,
+            load_format,
+            pause_generation_mode,
+        )
+        torch.cuda.synchronize()
+
     time_begin_update = time.perf_counter()
     if backend == "Engine":
         engine.update_weights_from_distributed(