refactor teacher server

[hjh0119]

No description provided.

[gemini-code-assist]

Code Review

This pull request refactors the Generalized Knowledge Distillation (GKD) trainer to fetch teacher logprobs using a new inference-only client (VLLMInferClient) and a dedicated /infer/ endpoint, while also lazy-loading transformers imports across utility files to optimize startup time. The code review identified several critical issues in the new implementation: self.teacher_client is only initialized on the main process, which will cause crashes on non-zero ranks during evaluation; incorrect tuple unpacking and an undefined variable encoded_chunkbatch in _assemble_topk_for_chunk will raise ValueError and NameError; parse_prompt_logprobs lacks safety checks for None values and padding, leading to potential AttributeError and ValueError crashes; and a potential KeyError exists when popping _teacher_raw from chunks.

[hjh0119]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the GKD (Generalized Knowledge Distillation) pipeline to fetch teacher logprobs more efficiently using a new /infer/ endpoint and VLLMInferClient, aligning both standard and Megatron-based trainers. The review feedback highlights critical issues, including a rank mismatch in the Megatron trainer where the teacher client is initialized on the last rank but called on rank 0, missing safety checks and padding in parse_prompt_logprobs that could cause shape mismatches or AttributeErrors, and potential KeyErrors when popping keys from batch dictionaries.