AllreduceCallback

class AllreduceCallback(reduce_method, group=WORLD, backend=None)[source]

Allreduce Callback with tensor fusion optimization.

Parameters

  • reduce_method (str) – the method to reduce gradiants. reduce_method should be “sum” or “mean”.

  • group( – attr:`.distributed.group.Group, optional): communication group. Default: WORLD.

  • backend (str, optional) – override distributed backend in allreduce. If backend is None, will use the backend set in dist.launcher. Default: None.

Examples

import megengine as mge
import megengine.autodiff as ad
import megengine.distributed as dist

gm = ad.GradManager()
gm.attach(linear_cls.parameters(), callbacks=[dist.make_allreduce_cb("sum")])

previous

make_allreduce_cb

next

megengine.distributed.helper.param_pack_split