From 98669e66e8035ec7b140ef2ca3991ef55e867b42 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Mon, 13 Jul 2026 12:26:15 +0000 Subject: [PATCH 1/6] feat(args): add --lora-ipc-weight-sync for LoRA-only IPC rollout sync --- miles/utils/arguments.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 0ada321d..fed11b5a 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1039,6 +1039,15 @@ def add_debug_arguments(parser): default=None, help="Override LoRA target modules. Default: per-model from TrainPipelineConfig.", ) + parser.add_argument( + "--lora-ipc-weight-sync", + action="store_true", + default=False, + help=( + "Sync only lora_A/lora_B to rollout via IPC with weight_update_mode=lora_merge " + "(requires matching sglang-d LoRAPipeline support)." + ), + ) parser.add_argument( "--diffusion-init-lora-weight", type=str, From 05eb25bcc85e80983d2f14332ef9c85e2fea4eb2 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Mon, 13 Jul 2026 12:26:18 +0000 Subject: [PATCH 2/6] feat(fsdp): add LoRA IPC weight updater with PEFT key mapping Sync only lora_A/lora_B tensors to rollout via weight_update_mode=lora_merge, avoiding full merged DiT weight gather and IPC transfer each step. --- .../diffusion_update_weight_utils.py | 171 +++++++++++++++++- 1 file changed, 166 insertions(+), 5 deletions(-) diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index dc6193c8..18dcbf61 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -1,8 +1,9 @@ import abc import logging import os +import re from argparse import Namespace -from collections.abc import Sequence +from collections.abc import Mapping, Sequence import ray import torch @@ -33,6 +34,70 @@ logger = logging.getLogger(__name__) +LORA_IPC_WEIGHT_UPDATE_MODE = "lora_merge" + + +class PeftLoRAKeyMapper: + """Map PEFT LoRA state-dict keys to sglang-d tensor names for IPC sync.""" + + _LORA_KEY_RE = re.compile(r"\.lora_([AB])(?:\.[^.]+)?(?:\.weight)?$") + _PEFT_PREFIX = "base_model.model." + + @classmethod + def is_lora_key(cls, name: str) -> bool: + return ".lora_A" in name or ".lora_B" in name + + @classmethod + def to_sgld_name(cls, name: str) -> str | None: + """Map a PEFT state-dict key to sglang-d LoRA tensor name.""" + if not cls.is_lora_key(name): + return None + + stripped = name + if stripped.startswith(cls._PEFT_PREFIX): + stripped = stripped[len(cls._PEFT_PREFIX) :] + + match = cls._LORA_KEY_RE.search(stripped) + if match is None: + return None + + layer_prefix = stripped[: match.start()] + ab = match.group(1) + return f"{layer_prefix}.lora_{ab}" + + @classmethod + def collect_sgld_names(cls, state_dict: Mapping[str, torch.Tensor]) -> set[str]: + names: set[str] = set() + for key in state_dict: + sgld_name = cls.to_sgld_name(key) + if sgld_name is not None: + names.add(sgld_name) + return names + + @classmethod + def collect_layer_prefixes(cls, state_dict: Mapping[str, torch.Tensor]) -> set[str]: + return {name.rsplit(".lora_", 1)[0] for name in cls.collect_sgld_names(state_dict)} + + @classmethod + def summarize_mapping( + cls, + state_dict: Mapping[str, torch.Tensor], + ) -> tuple[int, int, list[str], list[str]]: + """Return (num_tensors, num_layers, sample_layer_prefixes, unmapped_peft_keys).""" + sgld_names: list[str] = [] + unmapped: list[str] = [] + for key in state_dict: + if not cls.is_lora_key(key): + continue + sgld_name = cls.to_sgld_name(key) + if sgld_name is None: + unmapped.append(key) + else: + sgld_names.append(sgld_name) + layer_prefixes = {name.rsplit(".lora_", 1)[0] for name in sgld_names} + sample = sorted(layer_prefixes)[:5] + return len(sgld_names), len(layer_prefixes), sample, unmapped + class DiffusionUpdateWeight(abc.ABC): """Base updater used by diffusion training actors.""" @@ -88,12 +153,23 @@ def _update_component_weights(self, target_module: str, model: torch.nn.Module) self.wait_and_update_bucket_weights(bucket, target_module) del bucket - def wait_and_update_bucket_weights(self, bucket, target_module: str): + def wait_and_update_bucket_weights(self, bucket, target_module: str, weight_update_mode=None): bucket = [(name, param.wait()) if hasattr(param, "wait") else (name, param) for name, param in bucket] - self.update_bucket_weights(bucket, target_module, weight_version=self.weight_version) + self.update_bucket_weights( + bucket, + target_module, + weight_version=self.weight_version, + weight_update_mode=weight_update_mode, + ) @abc.abstractmethod - def update_bucket_weights(self, named_tensors, target_module: str, weight_version=None) -> None: + def update_bucket_weights( + self, + named_tensors, + target_module: str, + weight_version=None, + weight_update_mode: str | None = None, + ) -> None: pass @@ -123,7 +199,13 @@ def connect_rollout_engines( # Calculate TP rank within this SGLang engine group. self.tp_rank = dist.get_rank() - start_rank - def update_bucket_weights(self, named_tensors, target_module: str, weight_version=None) -> None: + def update_bucket_weights( + self, + named_tensors, + target_module: str, + weight_version=None, + weight_update_mode: str | None = None, + ) -> None: monkey_patch_torch_reductions() logger.info("Using flattened tensor bucket (diffusion updater, module=%s)", target_module) named_tensors_by_dtypes = {} @@ -173,6 +255,10 @@ def update_bucket_weights(self, named_tensors, target_module: str, weight_versio "target_modules": [target_module], "weight_version": str(weight_version), } + if weight_update_mode is not None: + kwargs["weight_update_mode"] = weight_update_mode + kwargs["lora_alpha"] = self.args.lora_alpha + kwargs["lora_rank"] = self.args.lora_rank ref = self._ipc_engine.update_weights_from_tensor.remote(**kwargs) ray.get(ref) @@ -322,3 +408,78 @@ def _verify_weight_sync(self, pairs: list[tuple[str, torch.Tensor]], target_modu all_equal = all(s == first_sum for _, s in engine_sums) pretty = " ".join(f"eng{idx}={s[:16] if isinstance(s, str) else s}" for idx, s in engine_sums) logger.warning(f"[weight_sync verify v{self.weight_version} cross-engine] " f"all_equal={all_equal} {pretty}") + + +class DiffusionUpdateWeightFromTensorLoRAIPC(DiffusionUpdateWeightFromTensor): + """Push only lora_A/lora_B tensors; rollout merges locally via weight_update_mode=lora_merge.""" + + def update_weights(self) -> None: + self.weight_version += 1 + for target_module, model in self.models.items(): + bucket: list[tuple[str, torch.Tensor]] = [] + bucket_size = 0 + num_lora_keys = 0 + unmapped_keys: list[str] = [] + + for name, param in model.state_dict().items(): + if not PeftLoRAKeyMapper.is_lora_key(name): + continue + sgld_name = PeftLoRAKeyMapper.to_sgld_name(name) + if sgld_name is None: + unmapped_keys.append(name) + continue + + param = param.cuda() + if isinstance(param, DTensor): + param = param.redistribute( + placements=[Replicate()] * param.device_mesh.ndim, + async_op=True, + ).to_local() + + sz = param.numel() * param.element_size() + if bucket and bucket_size + sz >= self.args.update_weight_buffer_size: + self.wait_and_update_bucket_weights( + bucket, + target_module, + weight_update_mode=LORA_IPC_WEIGHT_UPDATE_MODE, + ) + bucket, bucket_size = [], 0 + + bucket.append((sgld_name, param)) + bucket_size += sz + num_lora_keys += 1 + + if bucket: + self.wait_and_update_bucket_weights( + bucket, + target_module, + weight_update_mode=LORA_IPC_WEIGHT_UPDATE_MODE, + ) + + if self.weight_version <= 2 and dist.get_rank() == 0: + _, num_layers, sample_layers, _ = PeftLoRAKeyMapper.summarize_mapping(model.state_dict()) + logger.info( + "LoRA IPC weight sync v%s [%s]: pushed %d lora tensors, " "%d layer prefixes (unmapped=%d)", + self.weight_version, + target_module, + num_lora_keys, + num_layers, + len(unmapped_keys), + ) + if sample_layers: + logger.info( + "LoRA IPC [%s] sample layer prefixes: %s", + target_module, + sample_layers, + ) + if unmapped_keys: + logger.warning( + "LoRA IPC unmapped PEFT keys [%s] (first 5): %s", + target_module, + unmapped_keys[:5], + ) + if num_lora_keys == 0: + logger.error( + "LoRA IPC [%s]: no lora tensors found in training state_dict", + target_module, + ) From e08125507e618352586ae88e5cdb88782511b193 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Mon, 13 Jul 2026 12:26:20 +0000 Subject: [PATCH 3/6] feat(rollout): wire LoRA IPC sync into actor and sglang engine Select DiffusionUpdateWeightFromTensorLoRAIPC when enabled, and pass lora_merge metadata plus target modules to rollout weight updates. --- miles/backends/fsdp_utils/actor.py | 8 +++++++- .../sglang_diffusion_engine.py | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 561a9a4e..b2036d44 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -28,7 +28,11 @@ ) from . import checkpoint from .arguments import deterministic_capable_flash_fns -from .diffusion_update_weight_utils import DiffusionUpdateWeightFromTensor, DiffusionUpdateWeightFromTensorLoRA +from .diffusion_update_weight_utils import ( + DiffusionUpdateWeightFromTensor, + DiffusionUpdateWeightFromTensorLoRA, + DiffusionUpdateWeightFromTensorLoRAIPC, +) from .lr_scheduler import get_lr_scheduler from .parallel import create_fsdp_parallel_state @@ -185,6 +189,8 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty # sglang-d now supports /update_weights_from_tensor (PR #20464). if self.args.debug_train_only: self.weight_updater = None + elif self.args.use_lora and self.args.lora_ipc_weight_sync: + self.weight_updater = DiffusionUpdateWeightFromTensorLoRAIPC(self.args, self.models) elif self.args.use_lora: self.weight_updater = DiffusionUpdateWeightFromTensorLoRA(self.args, self.models) else: diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 8ae1ae9e..d4dc921f 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -246,6 +246,9 @@ def update_weights_from_tensor( load_format: str | None = None, target_modules: list[str] | None = None, weight_version: str | None = None, + weight_update_mode: str | None = None, + lora_alpha: int | None = None, + lora_rank: int | None = None, ): """ Update model weights from tensor data. The HTTP server will only post meta data, and the real weights will be copied directly from GPUs. @@ -261,6 +264,12 @@ def update_weights_from_tensor( payload["target_modules"] = target_modules if weight_version is not None: payload["weight_version"] = weight_version + if weight_update_mode is not None: + payload["weight_update_mode"] = weight_update_mode + if lora_alpha is not None: + payload["lora_alpha"] = lora_alpha + if lora_rank is not None: + payload["lora_rank"] = lora_rank return self._make_request( "update_weights_from_tensor", payload, @@ -308,6 +317,11 @@ def simulate_crash(self): def _compute_server_args(args, host, port, nccl_port): + from miles.backends.fsdp_utils.configs.train_pipeline_config import ( + get_train_pipeline_config_cls, + resolve_diffusion_model_family, + ) + # Only set fields SGL-D's ServerArgs actually accepts. GPU pinning is done # in `_init_normal` via CUDA_VISIBLE_DEVICES — SGL-D has no base_gpu_id arg. kwargs = { @@ -339,4 +353,9 @@ def _compute_server_args(args, host, port, nccl_port): if hasattr(args, f"sglang_{attr.name}") and attr.name not in kwargs: kwargs[attr.name] = getattr(args, f"sglang_{attr.name}") + if getattr(args, "use_lora", False) and getattr(args, "lora_ipc_weight_sync", False): + family = resolve_diffusion_model_family(args.diffusion_model) + train_pipeline_config = get_train_pipeline_config_cls(family)() + kwargs["lora_target_modules"] = args.lora_target_modules or train_pipeline_config.lora_target_modules + return kwargs From 9f762ab5c06385e908c91c35c2e5c360fb9edd7e Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Mon, 13 Jul 2026 12:26:23 +0000 Subject: [PATCH 4/6] scripts: enable LoRA IPC weight sync in SD3 and Qwen OCR runs --- scripts/run-diffusion-grpo-ocr-2gpu-flowgrpo-aligned.sh | 1 + scripts/run-diffusion-grpo-sd3-ocr-sglang.sh | 1 + 2 files changed, 2 insertions(+) diff --git a/scripts/run-diffusion-grpo-ocr-2gpu-flowgrpo-aligned.sh b/scripts/run-diffusion-grpo-ocr-2gpu-flowgrpo-aligned.sh index d4716eca..e2eac1f3 100755 --- a/scripts/run-diffusion-grpo-ocr-2gpu-flowgrpo-aligned.sh +++ b/scripts/run-diffusion-grpo-ocr-2gpu-flowgrpo-aligned.sh @@ -63,6 +63,7 @@ hf download --repo-type dataset rockdu/miles-diffusion-datasets \ --num-gpus-per-node 2 \ --colocate \ --use-lora \ + --lora-ipc-weight-sync \ --lora-rank 64 \ --lora-alpha 128 \ --diffusion-init-lora-weight gaussian \ diff --git a/scripts/run-diffusion-grpo-sd3-ocr-sglang.sh b/scripts/run-diffusion-grpo-sd3-ocr-sglang.sh index 84cdbf32..c191cbc3 100755 --- a/scripts/run-diffusion-grpo-sd3-ocr-sglang.sh +++ b/scripts/run-diffusion-grpo-sd3-ocr-sglang.sh @@ -111,6 +111,7 @@ python -u "${ROOT_DIR}/train_diffusion.py" \ --use-miles-router \ --sglang-server-concurrency 8 \ --use-lora \ + --lora-ipc-weight-sync \ --lora-rank 32 \ --lora-alpha 64 \ --diffusion-init-lora-weight gaussian \ From 6d2719ceb6e0e9fc0482c7f9d3fc32f7f6993cce Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Wed, 15 Jul 2026 03:04:44 +0000 Subject: [PATCH 5/6] fix(lora): resolve lora_target_modules in args and add mapper tests Fill per-model LoRA target modules during arg validation so rollout IPC sync does not re-resolve config at engine startup. Add CPU CI tests for PeftLoRAKeyMapper and server-args wiring. --- .../sglang_diffusion_engine.py | 9 +-- miles/utils/arguments.py | 11 +++ .../fsdp_utils/test_peft_lora_key_mapper.py | 68 +++++++++++++++++++ tests/fast/utils/test_lora_args.py | 34 ++++++++++ 4 files changed, 114 insertions(+), 8 deletions(-) create mode 100644 tests/fast/backends/fsdp_utils/test_peft_lora_key_mapper.py create mode 100644 tests/fast/utils/test_lora_args.py diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index ff33106e..ea1f5157 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -300,11 +300,6 @@ def simulate_crash(self): def _compute_server_args(args, host, port, nccl_port): - from miles.backends.fsdp_utils.configs.train_pipeline_config import ( - get_train_pipeline_config_cls, - resolve_diffusion_model_family, - ) - # Only set fields SGL-D's ServerArgs actually accepts. GPU pinning is done # in `_init_normal` via CUDA_VISIBLE_DEVICES — SGL-D has no base_gpu_id arg. kwargs = { @@ -332,9 +327,7 @@ def _compute_server_args(args, host, port, nccl_port): kwargs[attr.name] = getattr(args, f"sglang_{attr.name}") if getattr(args, "use_lora", False) and getattr(args, "lora_ipc_weight_sync", False): - family = resolve_diffusion_model_family(args.diffusion_model) - train_pipeline_config = get_train_pipeline_config_cls(family)() - kwargs["lora_target_modules"] = args.lora_target_modules or train_pipeline_config.lora_target_modules + kwargs["lora_target_modules"] = args.lora_target_modules # dit_precision / vae_precision are PipelineConfig fields, not ServerArgs, so forward them explicitly (only when changed from the class default, to avoid clobbering a subclass override). from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index fed11b5a..6268ffd4 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1369,6 +1369,17 @@ def miles_validate_args(args): if args.model_backend_path is None: args.model_backend_path = cfg_cls.model_backend_path cfg_cls.validate_args(args) + if args.use_lora and args.lora_target_modules is None: + args.lora_target_modules = list(cfg_cls.lora_target_modules) + + if getattr(args, "lora_ipc_weight_sync", False): + if not args.use_lora: + raise ValueError("--lora-ipc-weight-sync requires --use-lora") + if not args.lora_target_modules: + raise ValueError( + "--lora-ipc-weight-sync requires LoRA target modules; " + "set --diffusion-model (for per-model defaults) or --lora-target-modules." + ) if args.dump_details is not None: args.save_debug_rollout_data = f"{args.dump_details}/rollout_data/{{rollout_id}}.pt" diff --git a/tests/fast/backends/fsdp_utils/test_peft_lora_key_mapper.py b/tests/fast/backends/fsdp_utils/test_peft_lora_key_mapper.py new file mode 100644 index 00000000..d74aac0f --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_peft_lora_key_mapper.py @@ -0,0 +1,68 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[]) + +import pytest +import torch + +from miles.backends.fsdp_utils.diffusion_update_weight_utils import PeftLoRAKeyMapper + +_QWEN_A = "base_model.model.transformer_blocks.0.attn.to_q.lora_A.default.weight" +_QWEN_B = "base_model.model.transformer_blocks.0.attn.to_q.lora_B.default.weight" +_SD3_A = "base_model.model.transformer_blocks.1.attn.add_k_proj.lora_A.weight" +_MLP_B = "base_model.model.transformer_blocks.2.img_mlp.net.0.proj.lora_B.default.weight" + + +class TestPeftLoRAKeyMapper: + @pytest.mark.parametrize( + "key,expected", + [ + (_QWEN_A, "transformer_blocks.0.attn.to_q.lora_A"), + (_QWEN_B, "transformer_blocks.0.attn.to_q.lora_B"), + (_SD3_A, "transformer_blocks.1.attn.add_k_proj.lora_A"), + (_MLP_B, "transformer_blocks.2.img_mlp.net.0.proj.lora_B"), + ("transformer_blocks.0.attn.to_v.lora_A.default.weight", "transformer_blocks.0.attn.to_v.lora_A"), + ], + ) + def test_to_sgld_name_maps_peft_keys(self, key, expected): + assert PeftLoRAKeyMapper.to_sgld_name(key) == expected + + @pytest.mark.parametrize( + "key", + [ + "transformer_blocks.0.attn.to_q.weight", + "base_model.model.norm.weight", + "lora_A.default.weight", + ], + ) + def test_to_sgld_name_returns_none_for_non_lora_keys(self, key): + assert PeftLoRAKeyMapper.to_sgld_name(key) is None + + def test_is_lora_key(self): + assert PeftLoRAKeyMapper.is_lora_key(_QWEN_A) + assert PeftLoRAKeyMapper.is_lora_key(_QWEN_B) + assert not PeftLoRAKeyMapper.is_lora_key("transformer_blocks.0.attn.to_q.weight") + + def test_collect_sgld_names_and_layer_prefixes(self): + state_dict = { + _QWEN_A: torch.zeros(4, 8), + _QWEN_B: torch.zeros(8, 4), + "base_model.model.norm.weight": torch.zeros(8), + } + sgld_names = PeftLoRAKeyMapper.collect_sgld_names(state_dict) + assert sgld_names == { + "transformer_blocks.0.attn.to_q.lora_A", + "transformer_blocks.0.attn.to_q.lora_B", + } + assert PeftLoRAKeyMapper.collect_layer_prefixes(state_dict) == {"transformer_blocks.0.attn.to_q"} + + def test_summarize_mapping_reports_unmapped_lora_keys(self): + state_dict = { + _QWEN_A: torch.zeros(4, 8), + "base_model.model.weird.lora_A.default.weight.extra": torch.zeros(4, 8), + } + num_tensors, num_layers, sample_layers, unmapped = PeftLoRAKeyMapper.summarize_mapping(state_dict) + assert num_tensors == 1 + assert num_layers == 1 + assert sample_layers == ["transformer_blocks.0.attn.to_q"] + assert unmapped == ["base_model.model.weird.lora_A.default.weight.extra"] diff --git a/tests/fast/utils/test_lora_args.py b/tests/fast/utils/test_lora_args.py new file mode 100644 index 00000000..3f6092ea --- /dev/null +++ b/tests/fast/utils/test_lora_args.py @@ -0,0 +1,34 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[]) + +from argparse import Namespace + +from miles.backends.sglang_diffusion_utils.sglang_diffusion_engine import _compute_server_args + + +def _server_args(**overrides): + base = dict( + diffusion_model="Qwen/Qwen-Image", + diffusion_flow_shift=None, + rollout_num_gpus_per_engine=1, + sglang_sp_degree=None, + sglang_enable_cfg_parallel=False, + use_lora=True, + lora_ipc_weight_sync=True, + lora_target_modules=["to_q", "to_k"], + ) + base.update(overrides) + return Namespace(**base) + + +class TestLoRATargetModulesServerArgs: + def test_lora_ipc_uses_resolved_args(self): + args = _server_args() + kwargs = _compute_server_args(args, "127.0.0.1", 15000, 15001) + assert kwargs["lora_target_modules"] == ["to_q", "to_k"] + + def test_lora_ipc_omitted_when_disabled(self): + args = _server_args(lora_ipc_weight_sync=False) + kwargs = _compute_server_args(args, "127.0.0.1", 15000, 15001) + assert "lora_target_modules" not in kwargs From e6d7104482141e3307142f8270b4d60b5a1b29db Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Wed, 15 Jul 2026 07:12:00 +0000 Subject: [PATCH 6/6] fix: cpu ci --- tests/fast/utils/test_lora_args.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/fast/utils/test_lora_args.py b/tests/fast/utils/test_lora_args.py index 3f6092ea..2fc931e3 100644 --- a/tests/fast/utils/test_lora_args.py +++ b/tests/fast/utils/test_lora_args.py @@ -1,6 +1,6 @@ -from tests.ci.ci_register import register_cpu_ci +from tests.ci.ci_register import register_cuda_ci -register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[]) +register_cuda_ci(est_time=15, suite="stage-b-3-gpu-h200", labels=[]) from argparse import Namespace