colossalai.engine.ophooks.zero_hook

class colossalai.engine.ophooks.zero_hook.ZeroHook(shard_strategy, memstarts_collector, process_group=None)

A hook to process sharded param for ZeRO method.