TrainingWorker provides a Tinker-like API (https://thinkingmachines.ai/tinker/) as a RayWorkerGroup to a single controller. Currently, we only provide more coarse grained APIs, and do not provide exact APIs as Tinker does. But this can be added in the future.
| 46 | |
| 47 | |
| 48 | class TrainingWorker(Worker, DistProfilerExtension): |
| 49 | """ |
| 50 | TrainingWorker provides a Tinker-like API (https://thinkingmachines.ai/tinker/) as a RayWorkerGroup |
| 51 | to a single controller. Currently, we only provide more coarse grained APIs, |
| 52 | and do not provide exact APIs as Tinker does. But this can be added in the future. |
| 53 | """ |
| 54 | |
| 55 | def __init__(self, config: TrainingWorkerConfig): |
| 56 | Worker.__init__(self) |
| 57 | |
| 58 | from verl.workers.engine import BaseEngine, EngineRegistry |
| 59 | |
| 60 | initialize_global_process_group_ray(timeout_second=None) |
| 61 | |
| 62 | self.config = config |
| 63 | self.model_config = self.config.model_config |
| 64 | self.engine_config = self.config.engine_config |
| 65 | self.optimizer_config = self.config.optimizer_config |
| 66 | self.checkpoint_config = self.config.checkpoint_config |
| 67 | self.device_name = get_device_name() |
| 68 | |
| 69 | # we use the one defined in model |
| 70 | self.engine_config.use_remove_padding = self.model_config.use_remove_padding |
| 71 | |
| 72 | # TODO: add DistProfilerExtension |
| 73 | self.profiler_config = self.config.profiler_config |
| 74 | if self.profiler_config is not None: |
| 75 | self.profiler_tool_config = self.profiler_config.tool_config.get(self.profiler_config.tool, {}) |
| 76 | else: |
| 77 | self.profiler_tool_config = None |
| 78 | |
| 79 | DistProfilerExtension.__init__( |
| 80 | self, DistProfiler(rank=self.rank, config=self.profiler_config, tool_config=self.profiler_tool_config) |
| 81 | ) |
| 82 | |
| 83 | self.engine: BaseEngine = EngineRegistry.new( |
| 84 | model_type=self.config.model_type, |
| 85 | backend=self.engine_config.strategy, |
| 86 | model_config=self.model_config, |
| 87 | engine_config=self.engine_config, |
| 88 | optimizer_config=self.optimizer_config, |
| 89 | checkpoint_config=self.checkpoint_config, |
| 90 | ) |
| 91 | |
| 92 | # build dispatch info |
| 93 | self._register_dispatch_collect_info( |
| 94 | mesh_name="train", |
| 95 | dp_rank=self.engine.get_data_parallel_rank(), |
| 96 | is_collect=self.engine.is_mp_src_rank_with_outputs(), |
| 97 | ) |
| 98 | |
| 99 | self.flops_counter = FlopsCounter(self.model_config.hf_config) |
| 100 | |
| 101 | self.loss_fn = None |
| 102 | |
| 103 | @register(dispatch_mode=Dispatch.ONE_TO_ALL) |
| 104 | def to(self, device, model=True, optimizer=True, grad=True): |
| 105 | """Manual control of load/offload""" |
no outgoing calls
no test coverage detected