| 521 | |
| 522 | |
| 523 | class TensorParallelTrainer(BasicTrainer): |
| 524 | def __init__(self, policy, config, seed, run_dir, reference_model=None, rank=0, world_size=1): |
| 525 | """A trainer subclass that uses TensorParallel to shard the model across multiple GPUs. |
| 526 | |
| 527 | Based on https://github.com/BlackSamorez/tensor_parallel. Note sampling is extremely slow, |
| 528 | see https://github.com/BlackSamorez/tensor_parallel/issues/66. |
| 529 | """ |
| 530 | super().__init__(policy, config, seed, run_dir, reference_model, rank, world_size) |
| 531 | |
| 532 | rank0_print('Sharding policy...') |
| 533 | self.policy = tp.tensor_parallel(policy, sharded=True) |
| 534 | if config.loss.name in {'dpo', 'ipo'}: |
| 535 | rank0_print('Sharding reference model...') |
| 536 | self.reference_model = tp.tensor_parallel(reference_model, sharded=False) |
| 537 | |
| 538 | def save(self, output_dir=None, metrics=None): |
| 539 | """Save (unsharded) policy state to disk.""" |
| 540 | with tp.save_tensor_parallel(self.policy): |
| 541 | policy_state_dict = self.policy.state_dict() |
| 542 | |
| 543 | self.write_state_dict(self.example_counter, policy_state_dict, metrics, 'policy.pt', output_dir) |
| 544 | del policy_state_dict |
| 545 |
nothing calls this directly
no outgoing calls
no test coverage detected