[ROCM][RL] Shuffle Weight In-Place to Preserve Parameter Attributes

[zyzshishui]

Motivation

Several ROCm/aiter post-processing paths were replacing existing weight objects with newly constructed ones after shuffle_weight(...). That drops custom attributes attached to the original parameters, such as weight_loader, and breaks RL workflows that call load_weights() again after model initialization.

AttributeError: 'Parameter' object has no attribute 'weight_loader'

Modifications

Switche those paths to in-place .data updates to make original Parameter objects and their custom attributes remain intact.

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

[zyzshishui]

/tag-run-ci-label

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a5fa45e6ad

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[gemini-code-assist]

Code Review

This pull request refactors weight processing across several quantization schemes (FP8, MXFP4, Quark) to update torch.nn.Parameter data in-place using .data instead of re-assigning the parameter object. This change is intended to preserve parameter attributes and metadata. Feedback suggests that when transposing weights in-place, sharding attributes like input_dim and output_dim must be manually swapped to ensure correct behavior in workflows where weights might be re-loaded, such as reinforcement learning. Additionally, it is recommended to apply these in-place updates consistently across both aiter and non-aiter code paths and to include torch.cuda.empty_cache() for better memory management.

[gemini-code-assist]

Transposing the weight in-place without updating the sharding attributes (input_dim, output_dim) will cause incorrect sharding if load_weights is called again (e.g., in RL workflows), as the loader will use the old dimensions on the transposed shape. Both the aiter and non-aiter paths should be updated to handle this. Also, using .data for the input to shuffle_weight and adding torch.cuda.empty_cache() improves consistency and memory management.

Suggested change

	layer.weight.data = shuffle_weight(weight, (16, 16)).t()
	else:
	layer.weight = Parameter(weight.t(), requires_grad=False)
	layer.weight.data = shuffle_weight(weight.data, (16, 16)).t()
	else:
	layer.weight.data = weight.data.t()
	layer.weight.requires_grad = False
	# Update sharding attributes to reflect the transposed shape
	layer.weight.input_dim, layer.weight.output_dim = (
	layer.weight.output_dim,
	layer.weight.input_dim,
	)
	torch.cuda.empty_cache()

[HaiShaw]

@zyzshishui

[zyzshishui]

/rerun-failed-ci

[HaiShaw]

@kkHuang-amd please also review.

[HaiShaw]

@zyzshishui

[HaiShaw]

@amd-bot ci-status

[amd-bot]

@HaiShaw

CI Status for PR #21825

All 5 failures appear unrelated to your PR

Workflow	Job	Error	Related?	Log
PR Test	stage-b-test-1-gpu-large (13)	Fast-fail: skipping â�� root cause job(s): stage-b-test-1-gpu-large (11)	🟢 Unlikely -- The failure is in a 'check-stage-health' step which is an infrastructure/maintenance check unrelated to the quantization weight assignment changes in this PR.	View
PR Test	stage-b-test-1-gpu-large (11)	TimeoutError: The read operation timed out	🟢 Unlikely -- Without specific test error details linking to quantization weight loading, and given the PR only changes in-place assignment patterns on ROCm/aiter paths, this GPU test failure on what appears to be NVIDIA infrastructure is unlikely to be caused by these changes.	View
PR Test (NPU)	stage-b-test-16-npu-a3	Error: pod failed to come online with error: Error: Pod linux-aarch64-a3-16-pbjlp-runner-bdcqs-workflow is unhealthy with phase status Failed: {}	🟢 Unlikely -- The failure is a pod initialization error ('Pod is unhealthy with phase status Failed'), which is a Kubernetes/infrastructure issue completely unrelated to the PR's code changes.	View
PR Test (XPU)	build-and-test	AssertionError: Some test files in test suite do not exist on disk:	🟢 Unlikely -- This appears to be an Intel/NPU E2E bfloat16 test, and the PR only modifies ROCm-specific aiter weight shuffling paths that would not be exercised on this platform.	View
PR Test (Xeon)	build-test (all)	AssertionError: Some test files in test suite do not exist on disk:	🟢 Unlikely -- This appears to be an NPU/Intel build-test job, and the PR changes only affect ROCm/aiter-specific code paths for weight shuffling, making it unlikely to be related.	View

---

Generated by amd-bot using Claude API

[kkHuang-amd]

In process_weights_after_loading (per_channel), the _use_aiter branch updates layer.weight.data in place (preserves ModelWeightParameter / weight_loader), but the else branch still does layer.weight = Parameter(weight.t(), ...), which drops custom parameter types and attributes. If RL / load_weights() replay is meant to work only on ROCm+aiter, consider documenting that; otherwise the non-aiter path should also update .data in place (and set requires_grad=False if needed) for consistency.

[zyzshishui]

The corresponding feature in miles was just merged yesterday and quantization have not been enabled. The real issue I encountered was in unquant.py, I just found there are many other similar patterns and do those changes by the way. I’ve now also added in other non-aiter/rocm paths in the same schemes.

But I can also revert other changes and leave only the one I actually need in unquant.py, if you think it's better

[kkHuang-amd]

Same pattern as Quark: aiter path is in-place on the existing parameter, but the non-aiter branch replaces layer.weight with a new Parameter, which can still drop weight_loader and related metadata on non-ROCm runs. Worth aligning both branches if double load_weights() is a supported scenario there too.

[kkHuang-amd]

In-place layer.*.data = ... requires the RHS to match the existing parameter shape. Replacing Parameter(new_tensor) previously allowed shape changes if shuffling ever altered extents. Please confirm shuffled weights/scales always match the tensors allocated in create_weights for every code path; otherwise this can fail at runtime.

[zyzshishui]

In-place layer.*.data = ... requires the RHS to match the existing parameter shape.

I think that's not correct (? param.data = new_tensor is different from param.data.copy_(new_tensor), the latter requites match in shape. BTW I've called copy_or_rebind_param to handle this

[zyzshishui]

/rerun-failed-ci

[HaiShaw]

Only aiter path changed