The two ingredients of native amp (torch.cuda.amp.autocast and torch.cuda.amp.GradScaler) do not affect the model or optimizer in a stateful way.
autocast casts inputs to listed functions on the fly. The casted inputs are temporaries. They may end up stashed for backward, but that’s the longest they last. They never overwrite any model attributes.
GradScaler scales the loss and (either in an explicit call to scaler.unscale_(optimizer) or implicitly inside scaler.step(optimizer) unscales .grad attributes of the optimizer’s params in place, but it doesn’t overwrite or replace any optimizer attributes.
This is by design, to avoid the black box monkey-patching spaghetti that can take place in apex.amp.initialize. torch.cuda.amp has no equivalent to initialize, so just delete the line.
To toggle amp on or off based on run_config without needing to write divergent code, use autocast and GradScaler’s enabled= argument as shown here 3. There’s no equivalent of opt_levels with native amp, it’s either on or off, so you’ll have to decide how to interpret or change fp16_opt_level.