Getting parameters of embeddings （safe_get_local_fp32_param）and setting the weight of embeddings （safe_set_local_fp32_param） does not work (bug?).

Hello, we have recently been trying to use FGM to train LLM but have encountered difficulties about deepspeed. We would greatly appreciate your generous assistance.

Here are the describe of FGM:

For each x in embedding:
    1. Compute the forward loss and perform backward propagation to **obtain gradients** for x.
    2. Calculate r based on the gradient of the embedding matrix and **add it to the current embedding**, effectively x+r.
    3. Compute the forward loss for x+r and perform backward propagation to obtain the gradient for the adversarial updates. Accumulate this gradient with the one obtained in step 1.
    4. Restore the embedding to its original value obtained in step 1.
    5. Update the parameters based on the gradient obtained in step 3.

We have made our own attempts and consulted documentation, but after conducting experiments, we found that the loss did not change. Here are the key steps we followed:
(1) We used ‘safe_get_full_grad’ and ‘safe_get_local_grad’ to obtain the gradient values of the embedding layer after performing backward propagation on the model. We have confirmed that this approach is effective through experimentation.
(2) We used ‘safe_get_full_fp32_param’ and ‘safe_get_local_fp32_param’ to obtain the model’s weights, resulting in a one-dimensional tensor such as [158727].
(3) We attempted to modify the model’s weights directly using ‘safe_set_local_fp32_param’, but it seems to have had no effect, as the loss obtained after backward propagation, ‘loss2’, remained the same as ‘loss1’.

Here is the crucial code segment we used:

for batch in dataloader:
        [...]

        outputs = engine(**batch_with_device_assign)
        loss1 = outputs.loss.item()
        engine.backward(outputs.loss)

        if FGM_eps:
            # FGM : attack the weight of embed_tokens
            for n, lp in engine.named_parameters():
                if 'embed_tokens' in n:
                    # 1. backup the embedding weight
                    local_hp = safe_get_local_fp32_param(lp)
                    backup_local_hp = copy.deepcopy(local_hp.data)
                    backup_ds_tensor = copy.deepcopy(lp.ds_tensor)
                    #print(f"lp.data {lp.data}\n  lp.ds_tensor {lp.ds_tensor}\n local_hp {local_hp}\n  backup_local_hp {backup_local_hp}\n\n")    
                    #print(f"shape lp.data {lp.data.shape}  lp.ds_tensor {lp.ds_tensor.shape}  backup_local_hp {backup_local_hp.shape}\n\n")    
                    #print(f"origin lp.ds_tensor==backup_ds_tensor : { lp.ds_tensor.equal(backup_ds_tensor)}\n\n") # False, because of the dtype
                    
                    # 2. attack
                    hp_grad = safe_get_local_grad(lp)   # get the grad, work
                    norm = torch.norm(hp_grad)
                    if norm and not torch.isnan(norm):  
                        r_at = FGM_eps * hp_grad / norm
                        # print(f"r_at shape: {r_at.shape},local_hp.shape:{local_hp.shape}\n\n")
                        safe_set_local_fp32_param(lp, local_hp + r_at)    # althought the func can change the local_hp, but can not change ds_tensor, and the loss is no change.
                        new_local_hp = safe_get_local_fp32_param(lp)                    
                    
                    outputs = engine(**batch_with_device_assign)
                    
                    # print(f"forward lp.ds_tensor==backup_ds_tensor : { lp.ds_tensor.equal(backup_ds_tensor)}  new_local_hp==backup_local_hp : { new_local_hp.equal(backup_local_hp)}")     # True,  False
                    # print(f"loss1 {loss1}->loss2 {outputs.loss.item()}")    # equal loss
                    engine.backward(outputs.loss)
                    
                   # print(f"backward lp.ds_tensor==backup_ds_tensor : { lp.ds_tensor.equal(backup_ds_tensor)}  new_local_hp==backup_local_hp : { new_local_hp.equal(backup_local_hp)}")
                    
                    # 3. restore the embed_tokens
                    safe_set_local_fp32_param(lp, backup_local_hp)
                    print(f"restore lp.ds_tensor==backup_ds_tensor : { lp.ds_tensor.equal(backup_ds_tensor)}  new_local_hp==backup_local_hp : { new_local_hp.equal(backup_local_hp)}")     # True and True
                    break

        engine.step()

Even if ‘safe_set_local_fp32_param’ is successful in changing ‘local_hp’, ‘ds_tensor’ does not change. Furthermore, the loss obtained from the two calculations remains the same, which has left us wondering whether this is a bug.

We would appreciate your guidance on the following:

1. Have we used the functions incorrectly?
2. How can we perform these three functions: obtain gradient values, backup weights, and modify weights?

Thx.

ps: We try directly use engine.module.model.embed_tokens.weight.data.add_(r_at) to change the embedding weight, and it can get a different loss. Are there any risks with this approach?