Fix multi-GPU placement for GRACE, WISE, and ULTRAEDIT

[TengJiao33]

Summary

• keep GRACEAdapter on the same device as the edited layer
• move WISE activation masks to the activation-loss tensor device before multiplication
• move ULTRAEDIT / MALMEN hook indices to the layer input or gradient device before indexing
• keep existing hparams, public APIs, and editor call patterns unchanged

Validation

• git diff --check
• py_compile for the changed source files
• forced two-GPU Qwen2.5-1.5B smoke test for GRACE, WISE, and ULTRAEDIT
• tested both lower-layer and final-layer edits with the edited layer placed on each visible GPU
• result: 12/12 placement cases passed

This validation checks multi-GPU runtime/device placement. It does not claim changes to editing quality metrics.

[TengJiao33]

Validation details

I used a temporary forced two-GPU harness for the multi-GPU smoke test.

Setup:

• model: Qwen2.5-1.5B
• model split across two visible GPUs with a manual device_map
• input / hparams.device kept on the primary visible GPU
• edited layer tested on both visible GPUs

Methods tested:

• GRACE
• WISE
• ULTRAEDIT

Placement coverage per method:

• layer 0 on the secondary visible GPU
• layer 0 on the primary visible GPU
• final layer on the primary visible GPU
• final layer on the secondary visible GPU

Result:

• 12/12 placement cases passed