Skip to content

Create solstice.reduce() and solstice.gather() #2

@Charl-AI

Description

@Charl-AI
  • solsitce.reduce(): this should be a function reducing a list of pytrees to a single pytree. This should be a single dispatch generic function which calls .merge() on any Metrics objects (a bit like the CLU implementation) and has customisable behaviour for other types of pytrees (e.g. mean leaves, sum leaves etc...). Could add an implementation for model parallel training ???

  • solstice.gather(): possibly unnecessary. Utility for reducing and gathering sharded pytrees across devices. This will be most useful for accumulating Metrics, but maybe also for model parallel training ??

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions