diff --git a/rabit-learn/linear/Makefile b/rabit-learn/linear/Makefile
index 53102b8e9..ee76b03ce 100644
--- a/rabit-learn/linear/Makefile
+++ b/rabit-learn/linear/Makefile
@@ -18,4 +18,4 @@ linear.o: linear.cc ../../src/*.h linear.h ../solver/*.h
 # dependenies here
 linear.rabit: linear.o lib
 linear.mock: linear.o lib
-test.rabit: test_load.cc lib
+
diff --git a/rabit-learn/linear/test_load.cc b/rabit-learn/linear/test_load.cc
deleted file mode 100644
index dc486d6e5..000000000
--- a/rabit-learn/linear/test_load.cc
+++ /dev/null
@@ -1,33 +0,0 @@
-#include <rabit.h>
-#include "../io/io.h"
-
-int main(int argc, char *argv[]) {
-  using namespace rabit::io;
-  if (argc < 4) {
-    // intialize rabit engine
-    rabit::Init(argc, argv);
-    if (rabit::GetRank() == 0) {
-      rabit::TrackerPrintf("Usage: <data_in> npart rank\n");
-    }
-    rabit::Finalize();
-    return 0;
-  }
-  rabit::Init(argc, argv);
-  int n = 0;
-  InputSplit *in = CreateInputSplit(argv[1],
-                                    atoi(argv[2]),
-                                    atoi(argv[3]));
-  std::string line;
-  while (in->NextLine(&line)) {
-    if (n % 100 == 0) {
-      rabit::TrackerPrintf("[%d] finishes loading %d lines\n",
-                           rabit::GetRank(), n);
-    }
-    n++;
-  }
-  delete in;
-  rabit::TrackerPrintf("[%d] ALL finishes loading %d lines\n",
-                       rabit::GetRank(), n);
-  rabit::Finalize();
-  return 0;
-}
diff --git a/yarn/run_hdfs_prog.py b/yarn/run_hdfs_prog.py
index 65fc6104f..d3962bfa6 100755
--- a/yarn/run_hdfs_prog.py
+++ b/yarn/run_hdfs_prog.py
@@ -33,6 +33,13 @@ lpath.append('%s/jre/lib/amd64/server' % java_home)
 
 env = os.environ.copy()
 env['CLASSPATH'] = '${CLASSPATH}:' + (':'.join(cpath))
+
+# setup hdfs options
+if 'rabit_hdfs_opts' in env:
+    env['LIBHDFS_OPTS'] = env['rabit_hdfs_opts']
+elif 'LIBHDFS_OPTS' not in env:
+    env['LIBHDFS_OPTS'] = '--Xmx128m'
+
 env['LD_LIBRARY_PATH'] = '${LD_LIBRARY_PATH}:' + (':'.join(lpath)) 
 ret = subprocess.call(args = sys.argv[1:], env = env)
 sys.exit(ret)
diff --git a/yarn/src/org/apache/hadoop/yarn/rabit/ApplicationMaster.java b/yarn/src/org/apache/hadoop/yarn/rabit/ApplicationMaster.java
index a568b4f49..f75f4ae6d 100644
--- a/yarn/src/org/apache/hadoop/yarn/rabit/ApplicationMaster.java
+++ b/yarn/src/org/apache/hadoop/yarn/rabit/ApplicationMaster.java
@@ -60,7 +60,7 @@ public class ApplicationMaster {
     // total number of tasks
     private int numTasks = 1;
     // maximum number of attempts to try in each task
-    private int maxNumAttempt = 1;
+    private int maxNumAttempt = 3;
     // command to launch
     private String command = "";
 
@@ -340,6 +340,9 @@ public class ApplicationMaster {
             if (e.getKey().startsWith("rabit_")) {
                 env.put(e.getKey(), e.getValue());
             }
+            if (e.getKey() == "LIBHDFS_OPTS") {
+                env.put(e.getKey(), e.getValue());
+            }
         }
         env.put("rabit_task_id", String.valueOf(task.taskId));
         env.put("rabit_num_trial", String.valueOf(task.attemptCounter));
@@ -480,7 +483,7 @@ public class ApplicationMaster {
                 } catch (Exception e) {
                 }
                 LOG.info("[Rabit] Task " + r.taskId + " exited with status "
-                        + exstatus);
+                         + exstatus + " Diagnostics:"+ s.getDiagnostics());
                 failed.add(s.getContainerId());
             }
         }
diff --git a/yarn/src/org/apache/hadoop/yarn/rabit/Client.java b/yarn/src/org/apache/hadoop/yarn/rabit/Client.java
index 63f92d8be..7c6de9fb3 100644
--- a/yarn/src/org/apache/hadoop/yarn/rabit/Client.java
+++ b/yarn/src/org/apache/hadoop/yarn/rabit/Client.java
@@ -136,6 +136,9 @@ public class Client {
             if (e.getKey().startsWith("rabit_")) {
                 env.put(e.getKey(), e.getValue());
             }
+            if (e.getKey() == "LIBHDFS_OPTS") {
+                env.put(e.getKey(), e.getValue());
+            }
         }
         LOG.debug(env);
         return env;