MCPcopy Create free account
hub / github.com/deepspeedai/DeepSpeed / zero3_linear_wrap

Function zero3_linear_wrap

deepspeed/runtime/zero/linear.py:148–149  ·  view source on GitHub ↗
(input, weight, bias=None)

Source from the content-addressed store, hash-verified

146
147
148def zero3_linear_wrap(input, weight, bias=None):
149 return LinearFunctionForZeroStage3.apply(input, weight, bias)
150
151
152class LinearModuleForZeroStage3(Module):

Callers 6

test_grad_and_value_over_patched_functional_linearMethod · 0.90
loss_fnMethod · 0.90
test_unbatched_backward_matches_referenceMethod · 0.90
_run_forward_backwardMethod · 0.90
test_no_autocast_leak_into_backwardMethod · 0.90
test_setup_context_stores_autocast_attrsMethod · 0.90

Calls 1

applyMethod · 0.45

Tested by 6

test_grad_and_value_over_patched_functional_linearMethod · 0.72
loss_fnMethod · 0.72
test_unbatched_backward_matches_referenceMethod · 0.72
_run_forward_backwardMethod · 0.72
test_no_autocast_leak_into_backwardMethod · 0.72
test_setup_context_stores_autocast_attrsMethod · 0.72