[rollout] fix: remove dtype cast (#5117)

### What does this PR do?

- As title and comment

### 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`, `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.
  - 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`.

---------

Co-authored-by: wuxibin <wuxibin@bytedance.com>

Author: vermouth1992

verl/checkpoint_engine/hccl_checkpoint_engine.py

@@ -246,9 +246,6 @@ async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None,
         bucket_meta: dict[str, TensorMeta] = {}
         offset = 0
         for name, weight in weights:
-            # model parameters are in fp32 full precsion
-            weight = weight.to(self.rollout_dtype)
-
             # fill the tensor bucket
             if offset + weight.nbytes > self.bucket_size:
                 torch.npu.synchronize()

verl/checkpoint_engine/nccl_checkpoint_engine.py

@@ -240,9 +240,6 @@ async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None,
         bucket_meta: dict[str, TensorMeta] = {}
         offset = 0
         for name, weight in weights:
-            # model parameters are in fp32 full precsion
-            weight = weight.to(self.rollout_dtype)
-
             # fill the tensor bucket
             if offset + weight.nbytes > self.bucket_size:
                 torch.cuda.synchronize()

verl/checkpoint_engine/nixl_checkpoint_engine.py

@@ -385,9 +385,6 @@ async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None,
         bucket_meta: dict[str, TensorMeta] = {}
         offset = 0
         for name, weight in weights:
-            # model parameters are in fp32 full precision
-            weight = weight.to(self.rollout_dtype)
-
             # fill the tensor bucket
             if offset + weight.nbytes > self.bucket_size:
                 torch.cuda.synchronize()

verl/workers/engine/fsdp/transformer_impl.py

@@ -664,8 +664,14 @@ def get_per_tensor_param(self, layered_summon=False, base_sync_done=False):
             per_tensor_param = params
         else:
             device = get_device_id()  # used when fsdp2 set cpu_offload_policy
+            # TODO: cast fp32 to bf16 to reduce weight sync overhead, need more fine-grained control, e.g MoE gate
             per_tensor_param = (
-                (name, param.to(device, non_blocking=True).full_tensor() if isinstance(param, DTensor) else param)
+                (
+                    name,
+                    param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True)
+                    if isinstance(param, DTensor)
+                    else param,
+                )
                 for name, param in params.items()
             )
         return per_tensor_param, peft_config

verl/workers/rollout/vllm_rollout/vllm_rollout.py

@@ -42,7 +42,6 @@
 from verl import DataProto
 from verl.third_party.vllm import VLLM_SLEEP_LEVEL, get_version
 from verl.utils.device import get_device_id, get_device_name, get_torch_device, is_support_ipc
-from verl.utils.torch_dtypes import PrecisionType
 from verl.workers.config import HFModelConfig, RolloutConfig
 from verl.workers.rollout.base import BaseRollout
 from verl.workers.rollout.utils import ensure_async_iterator
@@ -190,10 +189,14 @@ async def update_weights(self, weights: Generator[tuple[str, torch.Tensor], None
         # send bucket weights
         offset = 0
         bucket_meta: dict[str, TensorMetadata] = {}
-        dtype = PrecisionType.to_dtype(self.config.dtype)
+        # dtype = PrecisionType.to_dtype(self.config.dtype)
         async for name, weight in ensure_async_iterator(weights):
             # model parameters are in fp32 full precision
-            weight = weight.to(dtype, non_blocking=True)
+            # (vermouth1992) we should not force cast weight here because some parameters
+            # (such as moe gate) have to keep fp32 precision. If a weight is bf16 in the rollout side,
+            # the rollout should automatically cast on demand. However, this would incur a higher weight
+            # transfer volume.
+            # weight = weight.to(dtype, non_blocking=True)
 
             # fill the tensor bucket
             if offset + weight.nbytes > bucket_size: