()
| 53 | |
| 54 | |
| 55 | def test(): |
| 56 | # construct model |
| 57 | ray.init() |
| 58 | |
| 59 | # create 2 workers, each hold a GPU |
| 60 | resource_pool = RayResourcePool([2], use_gpu=True, name_prefix="a") |
| 61 | |
| 62 | class_with_args = RayClassWithInitArgs(cls=ModelActor) |
| 63 | shard_wg = RayWorkerGroup(resource_pool, class_with_args, device_name=get_device_name()) |
| 64 | |
| 65 | test_bs = 8 |
| 66 | test_proto = DataProto( |
| 67 | TensorDict( |
| 68 | { |
| 69 | "sequence_ids": torch.ones([test_bs, 2048], dtype=torch.int64), |
| 70 | }, |
| 71 | batch_size=test_bs, |
| 72 | ), |
| 73 | meta_info={"query_length": 1536}, |
| 74 | ) |
| 75 | |
| 76 | # Sharding among different ranks |
| 77 | ret_proto1 = shard_wg.execute_with_func_generator(get_aux_metrics, test_proto) |
| 78 | |
| 79 | # compare execute on driver |
| 80 | hs = HackSelf() |
| 81 | ret_proto2 = get_aux_metrics(hs, test_proto) |
| 82 | |
| 83 | torch.testing.assert_close(ret_proto1.batch["decode_count"], ret_proto2.batch["decode_count"]) |
| 84 | |
| 85 | ray.shutdown() |
nothing calls this directly
no test coverage detected