feat: multi-lora training

[mathewjhan]

Summary

Allow training on multiple loras (all-linear, excluding expert) per step in Miles using megatron-bridge (related PR)

Feature

• Train multiple LoRAs in a single training step (currently only colocated until [WIP][lora] support disaggregate model lora training #988 is merged)
• Use as a long running service, supporting online loading and unloading
• Use as normal training, stopping when no more LoRAs are left to train

Running

see: examples/multi_lora

For normal training (not as a service):

1. run examples/multi_lora/provision.sh
2. configure W&B credentials and settings in examples/multi_lora/single_run.sh
3. run examples/multi_lora/single_run.sh |& tee run.log

For multi-lora training as a long running service:

1. run examples/multi_lora/provision.sh
2. configure W&B credentials and settings in examples/multi_lora/start_service.sh
3. run examples/multi_lora/start_service.sh |& tee run.log in one shell
4. run examples/multi_lora/submit_schedule.sh in a separate shell

Model checkpoints and LoRA safetensors are saved in `examples/adapters/*/checkpoints.

Changes to existing code (backwards compatible)

• Support --custom-generate-state flag to allow users to define their own GenerateState
• Add GenerateState hooks to allow custom generate states to access lifecycle of rollout, used for rollout request tracking
• Add AdapterRef and RewardSpec to Sample type so individual samples can access their own reward functions and adapter names during rollout

Notes

• Currently doesn't checkpoint optimizer + scheduler state yet, but can be added later as future PR
• Dataset checkpoint loading per adapter doesn't fully work yet since data source API doesn't support per adapter loading yet
• MultiLoRA not applied to experts as of now due to more complex bookkeeping required (need to keep track of the [adapter index, routed experts] together)
• The sglang lora csgmv kernel has some problems with the CUDA graph, so by default we use triton for now
• Doesn't support load from *.bin/*.safetensors yet, can be added later as a future PR, only resume from a megatron checkpoint or train from scratch

Tests

• e2e Qwen3-4B test using 2 LoRA adapters trained on gsm8k and dapo_math
• MultiLoRAController tests
• AdapterConfig tests

[gemini-code-assist]

Code Review

This pull request introduces Multi-LoRA training, enabling the concurrent training of multiple LoRA adapters against a shared base model with slot-based hot swapping. Key components include a central MultiLoRAController for lifecycle management, a MultiLoRADataSource for interleaved sampling, and backend updates to Megatron and SGLang for weight synchronization. Review feedback identifies several high-severity issues: a bug in the adapter registration sequence that prevents successful checkpoint resumption, a missing import causing a NameError during weight updates, and logic errors in step tracking for inactive adapters. Additionally, improvements were suggested for the efficiency of the singleton meta-class, the robustness of checkpoint path parsing, and the implementation of port-based liveness checks for Ray cluster initialization.

[gemini-code-assist]

load_adapter is called before init_adapter_slot. In the megatron-bridge implementation, init_adapter_slot is responsible for allocating the parameters for the specific LoRA slot. If it is called after load_adapter, the loading process will fail to find the target parameters, resulting in zero tensors being loaded and triggering the assertion failure on line 372. This effectively breaks checkpoint resumption for multi-LoRA.

Suggested change

	if ckpt is None:
	logger.info(f"{log_prefix} no checkpoint under {ckpt_root}, starting from random init")
	else:
	state_dict = torch.load(ckpt, map_location="cpu", weights_only=True)
	loaded = load_adapter(model, config.slot, state_dict)
	assert loaded > 0, (
	f"{log_prefix} loaded 0 tensors from {ckpt} "
	f"(state_dict has {len(state_dict)} entries) — name mismatch?"
	)
	logger.info(f"{log_prefix} loaded from {ckpt} ({loaded} tensors)")
	init_adapter_slot(model, config.slot, rank=config.rank, alpha=config.alpha)
	init_adapter_slot(model, config.slot, rank=config.rank, alpha=config.alpha)
	if ckpt is None:
	logger.info(f"{log_prefix} no checkpoint under {ckpt_root}, starting from random init")
	else:
	state_dict = torch.load(ckpt, map_location="cpu", weights_only=True)
	loaded = load_adapter(model, config.slot, state_dict)
	assert loaded > 0, (
	f"{log_prefix} loaded 0 tensors from {ckpt} "
	f"(state_dict has {len(state_dict)} entries) — name mismatch?"
	)
	logger.info(f"{log_prefix} loaded from {ckpt} ({loaded} tensors)")

[gemini-code-assist]

slice_lora_to_rank is used on line 302 but is not imported within this function or at the top level of the file. This will cause a NameError when update_multi_lora_weights is invoked.

Suggested change

	def update_multi_lora_weights(self, adapter_configs: dict[str, dict], active_slots: set[int] | None = None) -> None:
	"""Sync multiple LoRA adapters. Pause/resume once, loop export+send per adapter."""
	from megatron.bridge.peft.multi_lora_layers import expose_adapter_slot
	def update_multi_lora_weights(self, adapter_configs: dict[str, dict], active_slots: set[int] | None = None) -> None:
	"""Sync multiple LoRA adapters. Pause/resume once, loop export+send per adapter."""
	from megatron.bridge.peft.multi_lora_layers import expose_adapter_slot
	from miles.backends.megatron_utils.multi_lora_utils import slice_lora_to_rank

[gemini-code-assist]

Incrementing train_steps for all registered adapters in every training step is incorrect. In multi-LoRA training, an adapter might be skipped in a particular iteration (e.g., due to round-robin distribution or an empty data source). Since MultiLoRADataSource uses these steps to calculate data offsets for resumption and skipping, incrementing steps for adapters that did not contribute samples will cause them to skip data they never actually processed.

        # Increment the step count upon training completion.
        # TODO: Only increment for adapters that were actually included in the training batch
        # to avoid incorrect data skipping in MultiLoRADataSource.
        for name in self.train_steps.keys():
            self.train_steps[name] += 1

[gemini-code-assist]

When waiting for a server process like the Ray cluster to start, simply catching broad exceptions or checking process liveness is insufficient. According to repository guidelines, the check must verify that the server is actively listening for connections on its designated port (e.g., via a socket connection or HTTP request) to ensure it is truly ready.

Suggested change

	print("Waiting for Ray cluster to start...")
	except Exception as e:
	# TODO: Implement port-based liveness check per repository rules
	print(f"Waiting for Ray cluster to start... ({e})")

References

1. When waiting for a server process to start, verify it is listening on its port, not just process liveness.

[gemini-code-assist]

The logic for extracting the step number from directory names is fragile. It assumes that every directory starting with step_ has an integer as its second component. If a directory like step_latest or a system-generated folder exists, this will raise a ValueError or IndexError.

Suggested change

	get_step = lambda d: int(d.name.split("_")[1])
	get_step = lambda d: int(d.name.split("_")[1]) if len(d.name.split("_")) > 1 and d.name.split("_")[1].isdigit() else -1

[gemini-code-assist]

The implementation of SingletonMeta.__call__ is inefficient. It invokes cls.__new__ on every call, which in the case of GenerateState triggers dynamic class loading and uninitialized object creation even when the instance is already cached. The cache should be checked before calling __new__.

Suggested change

	def __call__(cls, *args, **kwargs):
	if cls not in cls._instances:
	instance = super().__call__(*args, **kwargs)
	cls._instances[cls] = instance
	return cls._instances[cls]
	# Note: this creates uninitialized instance copy of the
	# class even if it has been cached, for getting the type,
	# but the real instance will always be returned
	obj = cls.__new__(cls, *args, **kwargs)
	real_cls = type(obj)
	if real_cls not in cls._instances:
	obj.__init__(*args, **kwargs)
	cls._instances[real_cls] = obj
	return cls._instances[real_cls]
	def __call__(cls, *args, **kwargs):
	if cls in cls._instances:
	return cls._instances[cls]
	obj = cls.__new__(cls, *args, **kwargs)
	real_cls = type(obj)
	if real_cls not in cls._instances:
	obj.__init__(*args, **kwargs)
	cls._instances[real_cls] = obj
	if real_cls != cls:
	cls._instances[cls] = obj
	return cls._instances[real_cls]