[fix](kt-kernel): fix AMX BF16 Full-FT expert gradients

[Illumination111]

What does this PR do?

Fixes the AMX BF16 Full-FT path for CPU MoE expert base weights and removes a CPU-threading performance trap in SFT.

Root cause

• Full expert-weight gradients were sliced out by the TP configuration.
• TP children used a local input width where the global fused-weight stride was required.
• Token-major routing data was consumed as expert-major data.
• The down-projection dY snapshot needed by weight-gradient computation could be overwritten.
• Accelerate may inject OMP_NUM_THREADS=1 for a GPU launch, serializing CPU-side gradient accumulation, fused AdamW, zeroing, and requantization work.

Changes

• Complete the CPU BF16 expert chain from Parameter through C++ gradient buffers, autograd, optimizer update, and requantization.
• Correct TP slicing/strides, null checks, centralized zeroing, expert-major packed buffers, and retained down-projection dY.
• Add AMX BF16 weight-gradient kernels and NUMA-aware worker execution.
• Support Full/Hybrid/LoRA training, Qwen/GLM expert layouts, zero-storage placeholders, and requantization pointer reuse.
• Configure PyTorch/OpenMP threads from affinity-visible physical CPU cores. ACCELERATE_KT_OMP_NUM_THREADS has highest precedence; a user-provided OMP_NUM_THREADS > 1 remains respected. This host resolves to 96.
• Keep the existing fused AdamW optimizer. The experimental DeepSpeedCPUAdam hybrid was deliberately excluded: the controlled OMP=96 run made optimizer time 3.48x slower (6.804 s vs. 1.955 s) and increased step time from 20.405 s to 24.566 s.

Validation

• Full-FT structural run: 5/5 steps completed; 144/144 expert gradients were finite and nonzero; expert weights changed; backward fell from the earlier 754.5 s failure mode to 22.8 s.
• OMP=96 fused AdamW run: 20.405 s/step, 200.74 tokens/s, 9.239 s backward, 1.955 s optimizer; loss decreased from 0.7773 to 0.2754.
• pytest -q kt-kernel/test/per_commit/test_sft_omp_threads.py: 4 passed.
• Python compile checks and git diff --check: passed.

Known validation gap

The current end-to-end evidence covers finite/nonzero gradients, parameter updates, loss reduction, and stable performance. A full elementwise comparison against a PyTorch reference, including TP configurations, remains follow-up validation.

Before submitting

• I have read the contributing guidelines.
• I added focused tests for the new OpenMP thread-selection behavior.

[gemini-code-assist]

Code Review

This pull request implements Full Fine-Tuning (Full FT) support for KTransformers CPU/AMX MoE expert base weights (gate_proj, up_proj, down_proj). It establishes a complete data pipeline from Python SFT training to C++ AMX kernels, including autograd support, parallelized base weight gradient computation using AMX BF16 tiles, and an optimized re-quantization path that avoids full C++ object recreation. Additionally, it introduces GLM4 MoE architecture compatibility and replaces redundant expert weights in the model tree with zero-storage placeholders to prevent double-counting of parameters. The review feedback highlights a memory management issue in kt-kernel/python/sft/amx.py where using del old_moe fails to release the old C++ MOE object because self.moe still retains a reference. The reviewer suggests explicitly setting self.moe = None to ensure immediate garbage collection and prevent memory spikes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Using del old_moe only deletes the local variable reference, but self.moe still holds a reference to the old C++ MOE object. Consequently, the reference count of the old object does not drop to zero, and gc.collect() cannot release its memory before the new self.moe is instantiated. This causes both the old and new MOE objects to temporarily coexist in memory, which can lead to a significant memory spike and potential Out-Of-Memory (OOM) errors. Setting self.moe = None first will properly release the reference.

Suggested change

	# Release old C++ MOE object before creating a new one to avoid memory leak
	old_moe = getattr(self, "moe", None)
	if old_moe is not None:
	del old_moe
	import gc
	gc.collect()
	# Release old C++ MOE object before creating a new one to avoid memory leak
	if getattr(self, "moe", None) is not None:
	self.moe = None
	import gc
	gc.collect()

[gemini-code-assist]

Similar to the previous reference counting issue, del old_moe only deletes the local variable name, leaving self.moe still referencing the old C++ MOE object. To ensure the old object is immediately released and its memory can be reclaimed before loading the new weights, explicitly set self.moe = None instead.

Suggested change

	old_moe = getattr(self, "moe", None)
	if old_moe is not None:
	del old_moe
	if getattr(self, "moe", None) is not None:
	self.moe = None