wehub-resource-sync
147 lines
5.5 KiB
Python
147 lines
5.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""Conformance tests for the RolloutEngine interface.
|
|
|
|
Validates the dataclass invariants and exercises the interface against a
|
|
``FakeRollout`` so the contract is testable without GPUs or a model. The real
|
|
backends are tested manually with a launched training script (see README).
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from deepspeed.runtime.rollout import (
|
|
RolloutBatch,
|
|
RolloutEngine,
|
|
RolloutRequest,
|
|
SamplingConfig,
|
|
build_rollout,
|
|
)
|
|
|
|
# --- dataclass invariants ---------------------------------------------------
|
|
|
|
|
|
def test_rollout_request_validates_shapes():
|
|
with pytest.raises(ValueError, match="must be 2-D"):
|
|
RolloutRequest(prompt_ids=torch.zeros(8), prompt_attention_mask=torch.ones(8))
|
|
with pytest.raises(ValueError, match="does not match"):
|
|
RolloutRequest(prompt_ids=torch.zeros(2, 4, dtype=torch.long), prompt_attention_mask=torch.ones(2, 5))
|
|
|
|
|
|
def test_rollout_batch_validates_shapes():
|
|
with pytest.raises(ValueError, match="must be 2-D"):
|
|
RolloutBatch(input_ids=torch.zeros(8, dtype=torch.long),
|
|
attention_mask=torch.ones(8),
|
|
response_start_idx=torch.tensor([4]))
|
|
with pytest.raises(ValueError, match="does not match"):
|
|
RolloutBatch(input_ids=torch.zeros(2, 4, dtype=torch.long),
|
|
attention_mask=torch.ones(2, 5),
|
|
response_start_idx=torch.tensor([4, 4]))
|
|
with pytest.raises(ValueError, match="1-D of length"):
|
|
RolloutBatch(input_ids=torch.zeros(2, 4, dtype=torch.long),
|
|
attention_mask=torch.ones(2, 4),
|
|
response_start_idx=torch.tensor([4]))
|
|
|
|
|
|
def test_rollout_batch_accessors():
|
|
batch = RolloutBatch(
|
|
input_ids=torch.zeros(3, 12, dtype=torch.long),
|
|
attention_mask=torch.ones(3, 12),
|
|
response_start_idx=torch.tensor([4, 5, 6]),
|
|
)
|
|
assert batch.batch_size == 3
|
|
assert batch.seq_len == 12
|
|
|
|
|
|
def test_sampling_config_defaults():
|
|
cfg = SamplingConfig(max_new_tokens=32)
|
|
assert cfg.temperature == 1.0
|
|
assert cfg.top_p == 1.0
|
|
assert cfg.top_k == -1
|
|
assert cfg.n_samples_per_prompt == 1
|
|
|
|
|
|
# --- interface conformance via FakeRollout ---------------------------------
|
|
|
|
|
|
class FakeRollout(RolloutEngine):
|
|
"""Deterministic stub: appends ``[42] * max_new_tokens`` to each prompt."""
|
|
|
|
name = "fake"
|
|
|
|
def __init__(self, response_token: int = 42):
|
|
self.response_token = response_token
|
|
self.sync_calls: list = []
|
|
|
|
def generate(self, request: RolloutRequest, sampling: SamplingConfig) -> RolloutBatch:
|
|
B, T_p = request.prompt_ids.shape
|
|
n = sampling.n_samples_per_prompt
|
|
T_r = sampling.max_new_tokens
|
|
|
|
prompts_expanded = request.prompt_ids.repeat_interleave(n, dim=0)
|
|
attn_p_expanded = request.prompt_attention_mask.repeat_interleave(n, dim=0)
|
|
response = torch.full((B * n, T_r), self.response_token, dtype=request.prompt_ids.dtype)
|
|
response_attn = torch.ones((B * n, T_r), dtype=attn_p_expanded.dtype)
|
|
|
|
input_ids = torch.cat([prompts_expanded, response], dim=1)
|
|
attention_mask = torch.cat([attn_p_expanded, response_attn], dim=1)
|
|
response_start_idx = torch.full((B * n, ), T_p, dtype=torch.long)
|
|
return RolloutBatch(input_ids=input_ids, attention_mask=attention_mask, response_start_idx=response_start_idx)
|
|
|
|
def sync_weights(self, step: int) -> None:
|
|
self.sync_calls.append(step)
|
|
|
|
|
|
def test_fake_rollout_shape_basic():
|
|
fake = FakeRollout()
|
|
req = RolloutRequest(prompt_ids=torch.tensor([[1, 2, 3], [4, 5, 6]]),
|
|
prompt_attention_mask=torch.ones(2, 3, dtype=torch.long))
|
|
out = fake.generate(req, SamplingConfig(max_new_tokens=4))
|
|
assert out.input_ids.shape == (2, 7)
|
|
assert out.attention_mask.shape == (2, 7)
|
|
# With left-padded (fully real here) prompts of width 3, response begins
|
|
# at column 3 for every sample.
|
|
assert out.response_start_idx.tolist() == [3, 3]
|
|
|
|
|
|
def test_fake_rollout_with_n_samples():
|
|
fake = FakeRollout()
|
|
req = RolloutRequest(prompt_ids=torch.tensor([[1, 2], [3, 4]]),
|
|
prompt_attention_mask=torch.ones(2, 2, dtype=torch.long))
|
|
out = fake.generate(req, SamplingConfig(max_new_tokens=3, n_samples_per_prompt=4))
|
|
assert out.input_ids.shape == (8, 5)
|
|
assert out.response_start_idx.tolist() == [2] * 8
|
|
|
|
|
|
def test_fake_rollout_left_padded_prompts():
|
|
fake = FakeRollout()
|
|
# left-padded prompts: prompt B has only the last 2 positions real, but
|
|
# response_start_idx still equals the prompt column width T_p.
|
|
prompt_ids = torch.tensor([[1, 2, 3, 4], [0, 0, 5, 6]])
|
|
attn = torch.tensor([[1, 1, 1, 1], [0, 0, 1, 1]], dtype=torch.long)
|
|
req = RolloutRequest(prompt_ids=prompt_ids, prompt_attention_mask=attn)
|
|
out = fake.generate(req, SamplingConfig(max_new_tokens=2))
|
|
assert out.response_start_idx.tolist() == [4, 4]
|
|
|
|
|
|
def test_sync_records_steps():
|
|
fake = FakeRollout()
|
|
fake.sync_weights(0)
|
|
fake.sync_weights(5)
|
|
assert fake.sync_calls == [0, 5]
|
|
|
|
|
|
def test_engine_factory_unknown_raises():
|
|
from deepspeed.runtime.rollout.base import RolloutConfig
|
|
|
|
with pytest.raises(ValueError, match="Unknown rollout engine"):
|
|
build_rollout(RolloutConfig(engine="totally_made_up"))
|
|
|
|
|
|
def test_engine_factory_hybrid_requires_student_engine():
|
|
from deepspeed.runtime.rollout.base import RolloutConfig
|
|
|
|
with pytest.raises(ValueError, match="needs both"):
|
|
build_rollout(RolloutConfig(engine="hybrid_engine"))
|