[veomni] refactor: minor refactoring to ensure veomni engine compatibility with forward_only mode (#4889)

### What does this PR do?

Refactored the Vomni engine's initialize method into smaller
sub-methods, and enabling compatibility with forward_only mode.

### Checklist Before Starting

- [ ] Search for similar PRs. Paste at least one query link here: ...
- [ ] Format the PR title as `[{modules}] {type}: {description}` (This
will be checked by the CI)
- `{modules}` include `fsdp`, `megatron`, `sglang`, `vllm`, `rollout`,
`trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`,
`ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`,
`env`, `tool`, `ckpt`, `doc`, `data`, `cfg`, `reward`
- If this PR involves multiple modules, separate them with `,` like
`[megatron, fsdp, doc]`
  - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test`
- If this PR breaks any API (CLI arguments, config, function signature,
etc.), add `[BREAKING]` to the beginning of the title.
  - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching`

### Test

> For changes that can not be tested by CI (e.g., algorithm
implementation, new model support), validate by experiment(s) and show
results like training curve plots, evaluation results, etc.

### API and Usage Example

> Demonstrate how the API changes if any, and provide usage example(s)
if possible.

```python
# Add code snippet or script demonstrating how to use this
```

### Design & Code Changes

> Demonstrate the high-level design if this PR is complex, and list the
specific changes.

### Checklist Before Submitting

> [!IMPORTANT]
> Please check all the following items before requesting a review,
otherwise the reviewer might deprioritize this PR for review.

- [ ] Read the [Contribute
Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md).
- [ ] Apply [pre-commit
checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting):
`pre-commit install && pre-commit run --all-files --show-diff-on-failure
--color=always`
- [ ] Add / Update [the
documentation](https://github.com/volcengine/verl/tree/main/docs).
- [ ] Add unit or end-to-end test(s) to [the CI
workflow](https://github.com/volcengine/verl/tree/main/.github/workflows)
to cover all the code. If not feasible, explain why: ...
- [ ] Once your PR is ready for CI, send a message in [the `ci-request`
channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the
`verl` Slack
workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ).
(If not accessible, please try [the Feishu group
(飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).)
- [ ] If your PR is related to the `recipe` submodule, please also
update the reference to the submodule commit via `git submodule update
--remote` or `cd recipe && git pull origin main`.

Author: ji-huazhong

.github/PULL_REQUEST_TEMPLATE.md

@@ -6,7 +6,7 @@
 
 - [ ] Search for similar PRs. Paste at least one query link here: ...
 - [ ] Format the PR title as `[{modules}] {type}: {description}` (This will be checked by the CI)
-  - `{modules}` include `fsdp`, `megatron`, `sglang`, `vllm`, `rollout`, `trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`, `ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`, `env`, `tool`, `ckpt`, `doc`, `data`, `cfg`, `reward`
+  - `{modules}` include `fsdp`, `megatron`, `veomni`, `sglang`, `vllm`, `rollout`, `trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`, `ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`, `env`, `tool`, `ckpt`, `doc`, `data`, `cfg`, `reward`
   - If this PR involves multiple modules, separate them with `,` like `[megatron, fsdp, doc]`
   - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test`
   - If this PR breaks any API (CLI arguments, config, function signature, etc.), add `[BREAKING]` to the beginning of the title.

tests/special_sanity/check_pr_title.py

@@ -19,7 +19,7 @@
 pr_title = os.environ.get("PR_TITLE", "").strip()
 
 # Define rules
-allowed_modules = ["fsdp", "megatron", "sglang", "vllm", "rollout", "trainer"]
+allowed_modules = ["fsdp", "megatron", "veomni", "sglang", "vllm", "rollout", "trainer"]
 allowed_modules += ["tests", "training_utils", "recipe", "hardware", "deployment"]
 allowed_modules += ["ray", "worker", "single_controller", "misc", "docker", "ci"]
 allowed_modules += ["perf", "model", "algo", "env", "tool", "ckpt", "doc", "data", "cfg", "reward"]

verl/workers/engine/veomni/transformer_impl.py

@@ -31,12 +31,9 @@
 from verl.trainer.config import CheckpointConfig
 from verl.utils import tensordict_utils as tu
 from verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager
-from verl.utils.device import (
-    get_device_id,
-)
-from verl.utils.fsdp_utils import (
-    fsdp_version,
-)
+from verl.utils.device import get_device_id
+from verl.utils.fsdp_utils import fsdp_version
+from verl.utils.profiler import log_gpu_memory_usage
 from verl.workers.config import HFModelConfig, VeOmniEngineConfig, VeOmniOptimizerConfig
 from verl.workers.sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager
 
@@ -95,7 +92,6 @@ def __init__(
 
         self.use_remove_padding = self.model_config.use_remove_padding
 
-        # set FSDP offload params
         self._is_offload_param = self.engine_config.param_offload
         self._is_offload_optimizer = self.engine_config.optimizer_offload
         self._is_lora = self.model_config.lora_rank > 0
@@ -121,69 +117,100 @@ def __init__(
 
     def initialize(self):
         """
-        Build the model, optimizer, and learning rate scheduler under FSDP.
+        Build the model, optimizer, and learning rate scheduler under VeOmni.
 
         Applies device, dtype, and precision configurations, including mixed precision.
         Sets up checkpoint manager and FLOPs counter.
         """
+        self._build_model_optimizer()
 
-        self.module = build_foundation_model(
+        self.checkpoint_manager = FSDPCheckpointManager(
+            model=self.module,
+            optimizer=self.optimizer,
+            lr_scheduler=self.lr_scheduler,
+            processing_class=self.model_config.get_processor(),
+            checkpoint_config=self.checkpoint_config,
+        )
+
+        self.to(
+            device="cpu",
+            model=self._is_offload_param,
+            optimizer=self._is_offload_optimizer,
+            grad=self._is_offload_optimizer,
+        )
+
+        log_gpu_memory_usage("After offload model/optimizer/grad during init", logger=logger)
+
+    def _build_optimizer(self, module):
+        optimizer = build_optimizer(
+            module,
+            lr=self.optimizer_config.lr,
+            betas=self.optimizer_config.betas,
+            weight_decay=self.optimizer_config.weight_decay,
+            optimizer_type=self.optimizer_config.optimizer,
+        )
+        get_optimizer_pre_hook = getattr(module, "get_optimizer_pre_hook", None)
+        if get_optimizer_pre_hook is not None:
+            optimizer_pre_hook = get_optimizer_pre_hook(module, module.config, self.engine_config.data_parallel_mode)
+            optimizer.register_step_pre_hook(optimizer_pre_hook)
+
+        return optimizer
+
+    def _build_lr_scheduler(self, optimizer):
+        optim_config = self.optimizer_config
+        lr_scheduler = build_lr_scheduler(
+            optimizer,
+            train_steps=optim_config.total_training_steps,
+            lr=optim_config.lr,
+            lr_min=optim_config.lr_min,
+            lr_decay_style=optim_config.lr_scheduler_type,
+            lr_decay_ratio=optim_config.lr_decay_ratio,
+            lr_warmup_ratio=optim_config.lr_warmup_steps_ratio,
+            lr_start=optim_config.lr_start,
+        )
+
+        return lr_scheduler
+
+    def _build_model_optimizer(self):
+        # Load base model with specified configuration and dtype
+        module = build_foundation_model(
             config_path=self.model_config.hf_config_path,
             weights_path=self.model_config.path,
             torch_dtype="float32" if self.engine_config.mixed_precision else "bfloat16",
             attn_implementation=self.engine_config.attn_implementation,
             moe_implementation=self.engine_config.moe_implementation,
             init_device=self.engine_config.init_device,
         )
+        log_gpu_memory_usage("After load base model", logger=logger)
 
-        module_config = self.module.config
-
-        get_optimizer_pre_hook = getattr(self.module, "get_optimizer_pre_hook", None)
-        self.module = build_parallelize_model(
-            self.module,
+        # Applies parallel strategies to the model.
+        log_gpu_memory_usage("Before parallelize model", logger=logger)
+        module = build_parallelize_model(
+            module,
             init_device=self.engine_config.init_device,
             weights_path=self.model_config.path,
             enable_full_shard=self.engine_config.enable_full_shard,
             enable_mixed_precision=self.engine_config.mixed_precision,
             enable_gradient_checkpointing=self.model_config.enable_gradient_checkpointing,
             enable_fsdp_offload=self.engine_config.enable_fsdp_offload,
-            basic_modules=self.module._no_split_modules + self.engine_config.basic_modules,
+            basic_modules=module._no_split_modules + self.engine_config.basic_modules,
             enable_reentrant=self.engine_config.enable_reentrant,
             enable_forward_prefetch=self.engine_config.forward_prefetch,
         )
+        log_gpu_memory_usage("After parallelize model", logger=logger)
 
-        self.optimizer = build_optimizer(
-            self.module,
-            lr=self.optimizer_config.lr,
-            betas=self.optimizer_config.betas,
-            weight_decay=self.optimizer_config.weight_decay,
-            optimizer_type=self.optimizer_config.optimizer,
-        )
-        if get_optimizer_pre_hook is not None:
-            optimizer_pre_hook = get_optimizer_pre_hook(
-                self.module, module_config, self.engine_config.data_parallel_mode
-            )
-            self.optimizer.register_step_pre_hook(optimizer_pre_hook)
-
-        self.lr_scheduler = build_lr_scheduler(
-            self.optimizer,
-            train_steps=self.optimizer_config.total_training_steps,
-            lr=self.optimizer_config.lr,
-            lr_min=self.optimizer_config.lr_min,
-            lr_decay_style=self.optimizer_config.lr_scheduler_type,
-            lr_decay_ratio=self.optimizer_config.lr_decay_ratio,
-            lr_warmup_ratio=self.optimizer_config.lr_warmup_steps_ratio,
-            lr_start=self.optimizer_config.lr_start,
-        )
-
-        self.checkpoint_manager = FSDPCheckpointManager(
-            model=self.module,
-            optimizer=self.optimizer,
-            lr_scheduler=self.lr_scheduler,
-            processing_class=self.model_config.get_processor(),
-            checkpoint_contents=self.checkpoint_config,
-        )
+        if not self.engine_config.forward_only:
+            # Initialize optimizer with model parameters and config settings
+            optimizer = self._build_optimizer(module)
+            # Create learning rate scheduler with warmup and decay settings
+            lr_scheduler = self._build_lr_scheduler(optimizer)
+        else:
+            optimizer = None
+            lr_scheduler = None
 
+        self.module = module
+        self.optimizer = optimizer
+        self.lr_scheduler = lr_scheduler
         self.model_fwd_context, self.model_bwd_context = build_activation_offloading_context(
             self.model_config.enable_activation_offload,
             self.model_config.enable_gradient_checkpointing,