fix(export): prevent silent MLP linear collapse during HF export

[Surya-5555]

What does this PR do?

Type of change: Bug fix

This PR addresses the catastrophic drop in MMLU accuracy that occurs during Hugging Face export when the teacher is detached and the student's config is forced to activation_func = None (or an un-pickleable lambda/partial) for serialization.

Previously, when this sanitized config hit the _get_hidden_act export pipeline, it silently returned None or raised a silent failure. This caused the Multi-Layer Perceptron (MLP) to lose its non-linearity entirely, collapsing into a purely linear projection and destroying the model's accuracy.

Fixes:

1. Added strict validation to _get_hidden_act to raise an explicit ValueError if act_func is None, preventing the exporter from silently generating a non-activated linear model.
2. Added unrolling logic to safely extract the __name__ from functools.partial wrappers.
3. Added explicit checks to block anonymous <lambda> functions from passing silently.
4. Added 4 strict pytest unit tests in test_layer_utils.py to prevent future regressions.

Usage

N/A - This is an internal fix to the export pipeline safeguards.*

Testing

Ran local pytest unit tests verifying that ValueError exceptions are properly raised and halting the export pipeline for None, <lambda>, and unmapped inputs, while safely unwrapping functools.partial.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: N/A
• Did you get Claude approval on this PR?: N/A

Additional Information

Builds on top of the temporary fixes introduced in dkorzekwa/save_hf_checkpoint_for_every_val_iter_during_distill.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (3)

• main
• release/.*
• feature/.*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5d282d42-b256-4806-ae9f-cac3c5b7cc38

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[Surya-5555]

@AAnoosheh @kevalmorabia97
cc: @danielkorzekwa

This PR addresses the silent MLP linear collapse bug during Hugging Face export when activation_func evaluates to None (or an un-pickleable lambda/partial).

It introduces strict validation boundaries and unrolling logic for partial wrappers to ensure the exporter safely halts rather than generating a corrupted model.

Daniel requested your review on this fix so it can be merged into the distillation branch (targeting PR #1897). Please let me know if you have any feedback or require further modifications!

[danielkorzekwa]

Let's abandon this fix. We changed approach and will export all hf checkpoints once distillation is finished, see in MR: #1897