Function test_critic_engine

verl/tests/models/test_engine.py:250–350

(strategy)

@pytest.mark.parametrize("strategy", ["fsdp", "fsdp2"])
def test_critic_engine(strategy):
    device_count = torch.cuda.device_count()
    value_model_path = os.path.expanduser("~/models/test_model")
    language_model_path = get_test_language_model(device_count=device_count)
    create_value_model(language_model_path, value_model_path)

    torch.manual_seed(1)
    np.random.seed(1)

    ray.init()

    config = create_training_config(
        model_type="value_model", strategy=strategy, device_count=device_count, model=value_model_path
    )
    ray_cls_with_init = RayClassWithInitArgs(cls=ray.remote(TrainingWorker), config=config)
    resource_pool = RayResourcePool(process_on_nodes=[device_count])
    wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init)
    # init model
    wg.reset()

    batch_size = 8
    seqlen = 32

    response_length = seqlen // 2
    input_ids = torch.randint(0, config.model_config.hf_config.vocab_size, (batch_size, seqlen))
    attention_mask = create_random_mask(
        input_ids=input_ids, max_ratio_of_valid_token=0.8, max_ratio_of_left_padding=0.2, min_ratio_of_valid_token=0.6
    )
    position_ids = compute_position_id_with_mask(attention_mask)

    global_token_num = torch.sum(attention_mask, dim=-1).tolist()

    print(input_ids.float().mean(), attention_mask.float().mean())

    responses = input_ids[:, response_length:]
    response_mask = attention_mask[:, response_length:]

    assert torch.all(response_mask[:, 0] == 1)

    data = DataProto.from_single_dict(
        {
            "input_ids": input_ids,
            "prompts": input_ids[:, :response_length],
            "attention_mask": attention_mask,
            "position_ids": position_ids,
            "responses": responses,
            "response_mask": response_mask,
        },
        meta_info={"temperature": 1.0, "global_token_num": global_token_num, "compute_loss": False},
    )

    data_td = data.to_tensordict()
    data_td = left_right_2_no_padding(data_td)

    # eval
    output = wg.infer_batch(data_td)
    output = output.get()