[BUG] deepspeed amp seems to convert all input to specific dtype

Describe the bug
As title said, I made a minium snippet like this:

def compute_loss(inputs):
      input_ids = inputs.pop("input_ids")
      attention_mask = inputs.pop("attention_mask")

      valid_label_index_list = inputs.pop(
          "valid_label_index_list"
      ) 
      all_prob_supervised = inputs.pop("all_prob_supervised")
      all_prob_clm = inputs.pop("all_prob_clm")
      supervised_cnt = inputs.pop("supervised_cnt")
      clm_cnt = inputs.pop("clm_cnt")

      result = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )

If we set amp in deepspeed to 'bfloat16', all element in inputs will be convert from float32 to bfloat16.
While amp supported by huggingface trainer will not do this. So I am not sure if this is a feature of deepspeed or just a bug?
If this is a feature, how can we avoid this? (Since other input like clm_cnt in the case won't go into model)

System info (please complete the following information):

• GPU count and types 2xA100
• Python version 3.11
• deepspeed 0.14.2

Start by reproducing the minimal example with DeepSpeed AMP configured for bfloat16, using the stated 2xA100 environment if available, and compare it with Hugging Face Trainer AMP. Trace how AMP handles the full inputs mapping versus arguments passed to the model. Done means establishing whether this conversion is intended and, if not, defining behavior that preserves unrelated inputs such as clm_cnt.