Ray 分布式机器学习：利用Ray 进行大模型的数据处理、训练、推理和部署

Ray

生态及其他

|

231

pass

trainer = JaxTrainer(

...,

scaling_config=ScalingConfig(...),

datasets=dict(train=dataset),

)

from ray.train.data_parallel_trainer import DataParallelTrainer

class JaxTrainer(DataParallelTrainer):

最后，使用

Neptune

来记录和可视化

Tune Trials

，你可以定义一个

Logger-

Callback

，并将其传递到

Tuner

的运行配置中。要定义

LoggerCallback

，你需要

指定如何创建日志记录器（

setup

），在试验开始和结束时应该发生什么（

log_

trial_start

和

log_trial_end

），以及如何记录结果（

log_trial_result

）。如

果你已经实现了这样的类，例如

NeptuneCallback

，就可以像

11.1.2

节中使用

MLflowLogger Callback

一样使用它：

from ray.tune import logger, tuner

from ray.air.config import RunConfig

class NeptuneCallback(logger.LoggerCallback):

pass

tuner = tuner.Tuner(

trainer,

run_config=RunConfig(callbacks=[NeptuneCallback()]) ...