mirror of
https://github.com/volcengine/verl.git
synced 2025-10-20 21:53:50 +08:00
108 lines
3.3 KiB
Python
108 lines
3.3 KiB
Python
# Copyright 2024 Bytedance Ltd. and/or its affiliates
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
In this test, we instantiate a data parallel worker with 8 GPUs
|
|
"""
|
|
|
|
import ray
|
|
import tensordict
|
|
import torch
|
|
from codetiming import Timer
|
|
from torch import distributed as dist
|
|
|
|
from verl import DataProto
|
|
from verl.single_controller.base import Worker
|
|
from verl.single_controller.base.decorator import Dispatch, register
|
|
from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup
|
|
from verl.utils.ray_utils import parallel_put
|
|
|
|
|
|
@ray.remote
|
|
class DummyWorker(Worker):
|
|
def __init__(self):
|
|
super().__init__()
|
|
dist.init_process_group()
|
|
|
|
@register(dispatch_mode=Dispatch.DP_COMPUTE, blocking=False)
|
|
def do_nothing(self, data):
|
|
for key in data.batch.keys():
|
|
data.batch[key] += 1
|
|
if tensordict.__version__ >= "0.5.0":
|
|
data.batch = data.batch.consolidate()
|
|
return data
|
|
|
|
|
|
def test_data_transfer():
|
|
ray.init()
|
|
# construct resource pool
|
|
resource_pool = RayResourcePool([8])
|
|
cls_with_init = RayClassWithInitArgs(cls=DummyWorker)
|
|
# construct worker group
|
|
wg = RayWorkerGroup(resource_pool, cls_with_init)
|
|
|
|
# this is real dataset size
|
|
batch_size = 4096
|
|
seqlen = 32768
|
|
|
|
data_dict = {}
|
|
|
|
for i in range(2):
|
|
data_dict[str(i)] = torch.randint(0, 10000, (batch_size, seqlen))
|
|
|
|
data = DataProto.from_dict(tensors=data_dict)
|
|
|
|
print(data)
|
|
|
|
# we manually split data here and send to each worker
|
|
data_list = data.chunk(wg.world_size)
|
|
|
|
for i in range(wg.world_size):
|
|
# consolidate is necessary
|
|
if tensordict.__version__ >= "0.5.0":
|
|
data_list[i].batch = data_list[i].batch.consolidate()
|
|
|
|
with Timer(name="ray.pickle", initial_text=True):
|
|
for i in range(wg.world_size):
|
|
ray.cloudpickle.pickle.dumps(data_list[i])
|
|
|
|
with Timer(name="raw.pickle", initial_text=True):
|
|
import pickle
|
|
|
|
for i in range(wg.world_size):
|
|
pickle.dumps(data_list[i])
|
|
|
|
# we put in advance
|
|
with Timer(name="put", initial_text=True):
|
|
# takes around 40 seconds
|
|
data_list_ref = parallel_put(data_list)
|
|
# for i in range(wg.world_size):
|
|
# data_list[i] = ray.put(data_list[i])
|
|
|
|
with Timer(name="launch", initial_text=True):
|
|
output_ref = wg.do_nothing(data_list_ref)
|
|
|
|
with Timer(name="get", initial_text=True):
|
|
# takes around 40 seconds
|
|
output_lst = ray.get(output_ref)
|
|
|
|
for input_data, output_data in zip(data_list, output_lst):
|
|
for key in input_data.batch.keys():
|
|
assert torch.all(torch.eq(input_data.batch[key] + 1, output_data.batch[key])), (
|
|
input_data.batch[key],
|
|
output_data.batch[key],
|
|
key,
|
|
)
|
|
|
|
ray.shutdown()
|