Type checking in zero-3 to avoid silent bugs

This is a follow-up PR to the existing one: https://github.com/microsoft/DeepSpeed/pull/2127

The goal is to keep this PR open and investigate the issue in more detail while the PR above removes the annoying repeated printing of the warning.