Add dense Megatron LoRA-XS support

examples/train/lora_xs/run_megatron_lora_xs_qwen3-0.6b.sh

@@ -0,0 +1,73 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+set -x
+
+# Colocated GRPO for Qwen3-0.6B on GSM8K with Megatron and LoRA-XS.
+
+LORA_XS_INIT_METHOD="${LORA_XS_INIT_METHOD:-lora_xs}"
+LORA_RANK="${LORA_RANK:-32}"
+LORA_XS_INIT_DIR="${LORA_XS_INIT_DIR:-$HOME/lora_xs/qwen3_0.6b_${LORA_XS_INIT_METHOD}_r${LORA_RANK}}"
+SOURCE_MODEL="${SOURCE_MODEL:-Qwen/Qwen3-0.6B}"
+DEFAULT_POLICY_MODEL="$SOURCE_MODEL"
+if [[ "$LORA_XS_INIT_METHOD" == "pissa_xs" ]]; then
+  DEFAULT_POLICY_MODEL="$LORA_XS_INIT_DIR/residual_base"
+fi
+POLICY_MODEL="${POLICY_MODEL:-$DEFAULT_POLICY_MODEL}"
+DATA_DIR="${DATA_DIR:-$HOME/data/gsm8k}"
+NUM_GPUS="${NUM_GPUS:-8}"
+LOGGER="${LOGGER:-wandb}"
+LORA_ALPHA="${LORA_ALPHA:-$LORA_RANK}"
+
+uv run --isolated --extra megatron -m skyrl.train.entrypoints.main_base \
+  data.train_data="['$DATA_DIR/train.parquet']" \
+  data.val_data="['$DATA_DIR/validation.parquet']" \
+  trainer.algorithm.advantage_estimator=grpo \
+  trainer.policy.model.path="$POLICY_MODEL" \
+  trainer.ref.model.path="$SOURCE_MODEL" \
+  trainer.placement.colocate_all=true \
+  trainer.strategy=megatron \
+  trainer.placement.policy_num_gpus_per_node="$NUM_GPUS" \
+  trainer.placement.ref_num_gpus_per_node="$NUM_GPUS" \
+  generator.inference_engine.num_engines="$NUM_GPUS" \
+  generator.inference_engine.tensor_parallel_size=1 \
+  trainer.policy.megatron_config.tensor_model_parallel_size=1 \
+  trainer.policy.megatron_config.pipeline_model_parallel_size=1 \
+  trainer.policy.megatron_config.context_parallel_size=1 \
+  trainer.policy.megatron_config.lora_config.lora_type=lora_xs \
+  trainer.ref.megatron_config.tensor_model_parallel_size=1 \
+  trainer.ref.megatron_config.pipeline_model_parallel_size=1 \
+  trainer.ref.megatron_config.context_parallel_size=1 \
+  trainer.policy.model.lora.rank="$LORA_RANK" \
+  trainer.policy.model.lora.alpha="$LORA_ALPHA" \
+  trainer.policy.model.lora.target_modules=all-linear \
+  trainer.gradient_checkpointing=true \
+  trainer.remove_microbatch_padding=true \
+  trainer.epochs=20 \
+  trainer.eval_batch_size=1024 \
+  trainer.eval_before_train=false \
+  trainer.eval_interval=5 \
+  trainer.update_epochs_per_batch=1 \
+  trainer.train_batch_size=128 \
+  trainer.policy_mini_batch_size=64 \
+  trainer.micro_forward_batch_size_per_gpu=4 \
+  trainer.micro_train_batch_size_per_gpu=4 \
+  trainer.ckpt_interval=10 \
+  trainer.max_prompt_length=512 \
+  generator.sampling_params.max_generate_length=1024 \
+  trainer.policy.optimizer_config.lr=1.0e-5 \
+  trainer.algorithm.use_kl_loss=true \
+  generator.inference_engine.backend=vllm \
+  generator.inference_engine.run_engines_locally=true \
+  generator.inference_engine.weight_sync_backend=nccl \
+  generator.batched=true \
+  generator.n_samples_per_prompt=5 \
+  generator.inference_engine.gpu_memory_utilization=0.6 \
+  environment.env_class=gsm8k \
+  trainer.logger="$LOGGER" \
+  trainer.project_name=gsm8k_megatron \
+  trainer.run_name="gsm8k_megatron_qwen3_0.6b_lora_xs_r${LORA_RANK}_a${LORA_ALPHA}" \
+  trainer.resume_mode=from_path \
+  trainer.resume_path="$LORA_XS_INIT_DIR/global_step_0" \
+  trainer.ckpt_path="$HOME/ckpts/gsm8k_0.6b_lora_xs_ckpt" \
+  "$@"

examples/train/lora_xs/run_megatron_lora_xs_qwen3-0.6b_producer.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+BASE_MODEL="${BASE_MODEL:-Qwen/Qwen3-0.6B}"
+LORA_RANK="${LORA_RANK:-32}"
+LORA_XS_INIT_METHOD="${LORA_XS_INIT_METHOD:-lora_xs}"
+LORA_XS_INIT_DIR="${LORA_XS_INIT_DIR:-$HOME/lora_xs/qwen3_0.6b_${LORA_XS_INIT_METHOD}_r${LORA_RANK}}"
+DEFAULT_CONFIG_OVERRIDES='{"trainer.placement.policy_num_gpus_per_node":8,"trainer.policy.megatron_config.tensor_model_parallel_size":1,"trainer.policy.megatron_config.pipeline_model_parallel_size":1,"trainer.policy.megatron_config.context_parallel_size":1}'
+CONFIG_OVERRIDES="${CONFIG_OVERRIDES:-$DEFAULT_CONFIG_OVERRIDES}"
+
+uv run --isolated --extra megatron -m skyrl.train.entrypoints.lora_xs_init \
+  --base-model "$BASE_MODEL" \
+  --rank "$LORA_RANK" \
+  --init-method "$LORA_XS_INIT_METHOD" \
+  --output-dir "$LORA_XS_INIT_DIR" \
+  --config-overrides "$CONFIG_OVERRIDES"

skyrl/backends/skyrl_train/distributed/megatron/megatron_strategy.py

@@ -105,6 +105,7 @@ def _patched_update_fp32_params_by_new_state(self):
 _orig_load_parameter_state_from_dp_reshardable = DistributedOptimizer.load_parameter_state_from_dp_reshardable
 
 
+@torch.no_grad()
 def _patched_load_parameter_state_from_dp_reshardable(self, state_dict):
     """Wrapper around the original method that preserves the Adam step counter.
 

skyrl/backends/skyrl_train/distributed/megatron/optimizer.py

@@ -33,7 +33,7 @@
 
 
 def init_megatron_optim_config(
-    optim_config: Union[SkyRLOptimizerConfig, DictConfig], optimizer_config_kwargs: dict
+    optim_config: Union[SkyRLOptimizerConfig, DictConfig], optimizer_config_kwargs: dict, bf16: bool = True
 ) -> OptimizerConfig:
     adam_betas = getattr(optim_config, "adam_betas", (0.9, 0.999))
     optim_args = {
@@ -44,8 +44,9 @@ def init_megatron_optim_config(
         "weight_decay": getattr(optim_config, "weight_decay", 1e-2),
         "adam_beta1": float(adam_betas[0]),
         "adam_beta2": float(adam_betas[1]),
-        "bf16": True,
-        "params_dtype": torch.bfloat16,
+        "bf16": bf16,
+        "params_dtype": torch.bfloat16 if bf16 else torch.float32,
+        "store_param_remainders": bf16,
         "use_distributed_optimizer": True,
     }
     # YAML dtype overrides arrive as strings; Megatron expects torch.dtype.

skyrl/backends/skyrl_train/workers/megatron/lora_xs.py

@@ -0,0 +1,90 @@
+"""LoRA-XS adapters for Megatron training and PEFT export."""
+
+import contextlib
+
+import torch
+from megatron.bridge.peft.lora import LoRA
+from megatron.bridge.peft.lora_layers import LoRALinear, TEFusedLoRALinear
+from megatron.bridge.peft.utils import ParallelLinearAdapter
+from torch import nn
+
+LORA_XS_INIT_STD = 1e-5
+
+
+class LoRAXSCore(nn.Linear):
+    """Trainable square projection between frozen LoRA factors."""
+
+
+def configure_lora_xs_adapter(adapter: ParallelLinearAdapter) -> None:
+    """Freeze LoRA factors and insert a noise-initialized square core."""
+    if adapter.is_expert:
+        raise ValueError("LoRA-XS does not support expert adapters")
+    if not isinstance(adapter.activation, nn.Identity):
+        raise TypeError("LoRA-XS requires identity LoRA activation")
+
+    adapter.linear_in.weight.requires_grad_(False)
+    adapter.linear_out.weight.requires_grad_(False)
+    core = LoRAXSCore(
+        adapter.dim,
+        adapter.dim,
+        bias=False,
+        device=adapter.linear_in.weight.device,
+        dtype=adapter.linear_in.weight.dtype,
+    )
+    nn.init.normal_(core.weight, std=LORA_XS_INIT_STD)
+    adapter.activation = core
+
+
+class LoRAXS(LoRA):
+    """LoRA with frozen SVD factors and a trainable square core."""
+
+    def export_context(self, model_chunks):
+        """Materialize standard LoRA weights while exporting."""
+        return materialized_lora_xs_adapters(model_chunks)
+
+    def transform(self, module: nn.Module, name=None, prefix=None) -> nn.Module:
+        if isinstance(module, LoRALinear):
+            return module
+
+        transformed = super().transform(module, name, prefix)
+        if transformed is module:
+            return module
+        if isinstance(transformed, TEFusedLoRALinear):
+            transformed = LoRALinear(transformed.to_wrap, transformed.adapter)
+        if not isinstance(transformed, LoRALinear) or not isinstance(transformed.adapter, ParallelLinearAdapter):
+            raise TypeError("LoRA-XS supports only Megatron parallel linear adapters")
+
+        configure_lora_xs_adapter(transformed.adapter)
+        return transformed
+
+
+def _get_lora_xs_adapters(model_chunks) -> list[ParallelLinearAdapter]:
+    chunks = model_chunks if isinstance(model_chunks, (list, tuple)) else [model_chunks]
+    return [
+        module.adapter
+        for chunk in chunks
+        for module in chunk.modules()
+        if isinstance(module, LoRALinear)
+        and isinstance(module.adapter, ParallelLinearAdapter)
+        and isinstance(module.adapter.activation, LoRAXSCore)
+    ]
+
+
+@contextlib.contextmanager
+def materialized_lora_xs_adapters(model_chunks):
+    """Temporarily expose LoRA-XS adapters as standard two-factor LoRA."""
+    saved = []
+    try:
+        with torch.no_grad():
+            for adapter in _get_lora_xs_adapters(model_chunks):
+                core = adapter.activation
+                linear_out = adapter.linear_out.weight
+                saved.append((adapter, core, linear_out.detach().clone()))
+                linear_out.copy_(linear_out.float() @ core.weight.float())
+                adapter.activation = nn.Identity()
+        yield
+    finally:
+        with torch.no_grad():
+            for adapter, core, linear_out in saved:
+                adapter.linear_out.weight.copy_(linear_out)
+                adapter.activation = core

skyrl/backends/skyrl_train/workers/megatron/lora_xs_init_worker.py

@@ -0,0 +1,173 @@
+"""Megatron worker for offline LoRA-XS initialization."""
+
+import ray
+import torch
+from loguru import logger
+from megatron.bridge.peft.lora_layers import LoRALinear
+from megatron.bridge.peft.utils import ParallelLinearAdapter
+from megatron.core.distributed.distributed_data_parallel_config import (
+    DistributedDataParallelConfig,
+)
+
+from skyrl.backends.skyrl_train.workers.megatron.lora_xs import (
+    LORA_XS_INIT_STD,
+    LoRAXSCore,
+)
+from skyrl.backends.skyrl_train.workers.megatron.megatron_worker import (
+    MegatronPolicyWorkerBase,
+)
+from skyrl.backends.skyrl_train.workers.megatron.pissa_init_worker import (
+    _all_gather,
+    _shard,
+)
+from skyrl.train.config.config import get_config_as_dict
+
+
+def lora_xs_factors(weight: torch.Tensor, rank: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return the principal right vectors and singular-value-scaled left vectors."""
+    max_rank = min(weight.shape)
+    if rank > max_rank:
+        raise ValueError(f"LoRA-XS rank {rank} exceeds maximum rank {max_rank} for weight shape {tuple(weight.shape)}")
+
+    u, s, vh = torch.linalg.svd(weight.float(), full_matrices=False)
+    linear_in = (s[:rank].unsqueeze(1) * vh[:rank, :]).contiguous()
+    linear_out = u[:, :rank].contiguous()
+    return linear_in, linear_out
+
+
+def _synchronized_lora_xs_factors(weight: torch.Tensor, rank: int, group) -> tuple[torch.Tensor, torch.Tensor]:
+    if torch.distributed.get_rank(group) == 0:
+        linear_in, linear_out = lora_xs_factors(weight, rank)
+    else:
+        linear_in = torch.empty((rank, weight.shape[1]), dtype=torch.float32, device=weight.device)
+        linear_out = torch.empty((weight.shape[0], rank), dtype=torch.float32, device=weight.device)
+
+    torch.distributed.broadcast(linear_in, group=group, group_src=0)
+    torch.distributed.broadcast(linear_out, group=group, group_src=0)
+    return linear_in, linear_out
+
+
+@torch.no_grad()
+def _init_one_lora_xs_adapter(
+    base_linear,
+    adapter,
+    tp_size: int,
+    tp_rank: int,
+    tp_group,
+    use_residual_base: bool,
+) -> None:
+    base_weight = base_linear.weight
+    if base_weight.is_meta:
+        raise RuntimeError("LoRA-XS requires pretrained weights before adapter initialization")
+
+    base_shard_dim = 1 if adapter.input_is_parallel else 0
+    if tp_size == 1:
+        full_weight = base_weight.detach().float()
+        linear_in, linear_out = lora_xs_factors(full_weight, adapter.dim)
+    else:
+        full_weight = _all_gather(base_weight.detach().float(), base_shard_dim, tp_size, tp_group)
+        linear_in, linear_out = _synchronized_lora_xs_factors(full_weight, adapter.dim, tp_group)
+
+    linear_in_shard_dim = 1 if adapter.input_is_parallel else 0
+    dtype = adapter.linear_in.weight.dtype
+    adapter.linear_in.weight.copy_(_shard(linear_in, linear_in_shard_dim, tp_rank, tp_size).to(dtype))
+    adapter.linear_out.weight.copy_(_shard(linear_out, 0, tp_rank, tp_size).to(dtype))
+    adapter.activation.weight.normal_(std=LORA_XS_INIT_STD)
+    if use_residual_base:
+        adapter.activation.weight.add_(torch.eye(adapter.dim, device=base_weight.device, dtype=dtype))
+        scale = adapter.alpha / adapter.dim
+        residual = full_weight - scale * (linear_out @ linear_in)
+        base_weight.copy_(_shard(residual, base_shard_dim, tp_rank, tp_size).to(base_weight.dtype))
+
+
+def apply_lora_xs_init(model_chunks, use_residual_base: bool = False) -> None:
+    """Initialize dense LoRA-XS adapters from their base weights."""
+    import megatron.core.parallel_state as mpu
+
+    tp_size = mpu.get_tensor_model_parallel_world_size()
+    tp_rank = mpu.get_tensor_model_parallel_rank()
+    tp_group = mpu.get_tensor_model_parallel_group()
+    chunks = model_chunks if isinstance(model_chunks, (list, tuple)) else [model_chunks]
+    adapters = []
+    for chunk in chunks:
+        for module in chunk.modules():
+            if isinstance(module, LoRALinear) and isinstance(module.adapter, ParallelLinearAdapter):
+                if not isinstance(module.adapter.activation, LoRAXSCore):
+                    raise TypeError("LoRA-XS found a non-LoRA-XS parallel adapter")
+                adapters.append((module.to_wrap, module.adapter))
+
+    if not adapters:
+        raise ValueError("LoRA-XS found no supported adapters")
+    for base_linear, adapter in adapters:
+        _init_one_lora_xs_adapter(
+            base_linear,
+            adapter,
+            tp_size,
+            tp_rank,
+            tp_group,
+            use_residual_base,
+        )
+    logger.info(f"LoRA-XS: initialized {len(adapters)} adapter(s)")
+
+
+class LoRAXSInitWorkerBase(MegatronPolicyWorkerBase):
+    use_residual_base = False
+
+    def make_megatron_module(
+        self,
+        wrap_with_ddp=True,
+        ddp_config=None,
+        lora_config=None,
+        lora_type="lora_xs",
+        bf16=True,
+    ):
+        if lora_type != "lora_xs":
+            raise ValueError("LoRA-XS initialization requires lora_type='lora_xs'")
+        self.configure_lora(lora_config, lora_type)
+
+        def lora_pre_wrap_hook(model):
+            lora_model = self.lora_cls(model, training=True)
+            self.lora_cls.set_params_to_save(lora_model)
+            return lora_model
+
+        def lora_xs_pre_wrap_hook(model):
+            apply_lora_xs_init(model, self.use_residual_base)
+            return model
+
+        self.provider.register_pre_wrap_hook(lora_pre_wrap_hook)
+        self.provider.register_pre_wrap_hook(lora_xs_pre_wrap_hook)
+
+        resolved_ddp_config = DistributedDataParallelConfig()
+        if wrap_with_ddp:
+            resolved_ddp_config.use_distributed_optimizer = True
+        if ddp_config is not None:
+            for key, value in get_config_as_dict(ddp_config).items():
+                setattr(resolved_ddp_config, key, value)
+        return self.provider.provide_distributed_model(
+            ddp_config=resolved_ddp_config,
+            wrap_with_ddp=wrap_with_ddp,
+            bf16=bf16,
+        )
+
+    def save_residual_base(self, export_dir: str, tokenizer) -> None:
+        from skyrl.backends.skyrl_train.workers.megatron.pissa_init_worker import (
+            zeroed_adapters,
+        )
+
+        with zeroed_adapters(self.model.actor_module):
+            self.strategy.save_hf_model(
+                self.bridge,
+                self.model,
+                export_dir,
+                tokenizer=tokenizer,
+            )
+
+
+class PiSSAXSInitWorkerBase(LoRAXSInitWorkerBase):
+    """Megatron worker for residual-base PiSSA-XS initialization."""
+
+    use_residual_base = True
+
+
+LoRAXSInitWorker = ray.remote(num_gpus=1)(LoRAXSInitWorkerBase)
+PiSSAXSInitWorker = ray.remote(num_gpus=1)(PiSSAXSInitWorkerBase)