モデルの勾配パラメータを移動し、メモリ効率を改善
## 概要
モデル同期の際に勾配テンソルが推論モデルに移ってしまい、二重に勾配テンソルが存在していたためVRAMの使用量が増加していた。
勾配テンソルを移動する処理を記述することでメモリ使用量を改善した。
VRAM使用量は、Largeサイズで 22.5 GiB -> 19.98 GiBへ減少した。
<!-- 変更の目的 もしくは 関連する Issue 番号 -->
## 変更内容
<!-- ビューの変更がある場合はスクショによる比較などがあるとわかりやすい -->
## 影響範囲
若干学習時間が増える可能性があるが、動かしてみた時はあまり伸びているように感じなかった。
<!-- この関数を変更したのでこの機能にも影響がある、など -->
## Submit前の確認項目
<!-- PRをSubmitする前に確認する項目 -->
- [x] タイトルは一目でわかるようにし、説明文は PR を簡潔に説明するようにしましたか?
- [x] あなたの PR が、異なる変更を束ねたものではなく、ひとつのことだけを行うものであることを確認しましたか?
- [x] この PR で導入されたすべての変更点をリストアップしましたか?
- [x] PR を`make run`コマンドでローカルでテストしましたか?
## 補足
<!-- レビューをする際に見てほしい点、ローカル環境で試す際の注意点、など -->
合并状态:已合并 合并于 2025-02-19 关闭于 2025-02-19 0 条评论